From 7261238eede941ed8bba877159f9cff5fa8b5fad Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:18:39 +0000 Subject: [PATCH 001/162] feat(distributed): store the active carrier in a cluster row Add table cluster_carrier and a store for it. The row names the carrier that the whole cluster uses, an epoch that grows with every change, and the fields a later change of carrier needs: state, target and the carrier that is draining. Seed inserts the row only when it is absent, so replicas that start together agree on one writer. Transition is a compare-and-set on the epoch, so two writers cannot interleave. A check constraint keeps the table at one row. Nothing reads the row yet. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/carrier_store.go | 193 ++++++++++++++++ core/services/cluster/carrier_store_test.go | 236 ++++++++++++++++++++ core/services/cluster/cluster_suite_test.go | 13 ++ 3 files changed, 442 insertions(+) create mode 100644 core/services/cluster/carrier_store.go create mode 100644 core/services/cluster/carrier_store_test.go create mode 100644 core/services/cluster/cluster_suite_test.go diff --git a/core/services/cluster/carrier_store.go b/core/services/cluster/carrier_store.go new file mode 100644 index 000000000000..73f664703ce5 --- /dev/null +++ b/core/services/cluster/carrier_store.go @@ -0,0 +1,193 @@ +// Package cluster holds state that every frontend replica of a distributed +// deployment shares through the database. +package cluster + +import ( + "context" + "errors" + "fmt" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// Carrier names a transport that carries fan-out messages, work queues and +// control requests between frontends and workers. Exactly one carrier is +// active in a cluster at a time. +type Carrier string + +const ( + CarrierNATS Carrier = "nats" + CarrierTunnel Carrier = "tunnel" +) + +func (c Carrier) valid() bool { return c == CarrierNATS || c == CarrierTunnel } + +// State is where the cluster is in a change of carrier. +type State string + +const ( + // StateStable: one carrier is active and nothing is changing. + StateStable State = "stable" + // StatePrepare: replicas build the target carrier. The active carrier is + // unchanged. + StatePrepare State = "prepare" + // StateCommit: the target is now the active carrier and the previous one + // is still attached while it drains. + StateCommit State = "commit" +) + +func (s State) valid() bool { return s == StateStable || s == StatePrepare || s == StateCommit } + +var ( + // ErrNotSeeded means no replica has written the carrier row yet. + ErrNotSeeded = errors.New("cluster carrier row does not exist") + // ErrStaleEpoch means another writer changed the row after the caller read it. + ErrStaleEpoch = errors.New("cluster carrier epoch is stale") + // ErrInvalidCarrier means a carrier name is not one the cluster knows. + ErrInvalidCarrier = errors.New("unknown cluster carrier") + // ErrInvalidChange means a change does not describe a valid state. + ErrInvalidChange = errors.New("invalid cluster carrier change") +) + +// carrierRowID is the primary key of the only row. A check constraint keeps +// the table at one row. +const carrierRowID = 1 + +// CarrierRow is the one row of table cluster_carrier. Active is the carrier +// that publishes, enqueues and dials. Epoch grows with every change, so a +// reader can tell a newer row from an older one without comparing clocks. +type CarrierRow struct { + ID int `gorm:"primaryKey;check:cluster_carrier_single_row,id = 1" json:"-"` + Active Carrier `gorm:"size:16;not null" json:"active"` + Epoch int64 `gorm:"not null" json:"epoch"` + State State `gorm:"size:16;not null" json:"state"` + Target Carrier `gorm:"size:16;not null;default:''" json:"target,omitempty"` + Draining Carrier `gorm:"size:16;not null;default:''" json:"draining,omitempty"` + DrainingUntil *time.Time `json:"draining_until,omitempty"` + ChangedBy string `gorm:"size:255" json:"changed_by"` + // ChangedAt is for operators. Nothing decides on it, so replicas need no + // agreed clock. + ChangedAt time.Time `json:"changed_at"` + PrevEpoch int64 `json:"prev_epoch"` +} + +func (CarrierRow) TableName() string { return "cluster_carrier" } + +// Change is the next state of the row. Every field is written; a field left +// empty clears the column. +type Change struct { + Active Carrier + State State + Target Carrier + Draining Carrier + DrainingUntil *time.Time + By string +} + +func (c Change) validate() error { + if !c.Active.valid() { + return fmt.Errorf("%w: active %q: %w", ErrInvalidChange, c.Active, ErrInvalidCarrier) + } + if !c.State.valid() { + return fmt.Errorf("%w: state %q", ErrInvalidChange, c.State) + } + if c.Target != "" && !c.Target.valid() { + return fmt.Errorf("%w: target %q: %w", ErrInvalidChange, c.Target, ErrInvalidCarrier) + } + if c.Draining != "" && !c.Draining.valid() { + return fmt.Errorf("%w: draining %q: %w", ErrInvalidChange, c.Draining, ErrInvalidCarrier) + } + if c.State == StateStable && c.Target != "" { + return fmt.Errorf("%w: a stable state has no target", ErrInvalidChange) + } + if c.State != StateStable && c.Target == "" { + return fmt.Errorf("%w: state %q needs a target", ErrInvalidChange, c.State) + } + return nil +} + +// CarrierStore reads and writes the carrier row. +type CarrierStore struct { + db *gorm.DB +} + +// NewCarrierStore creates the table when it is missing. +func NewCarrierStore(db *gorm.DB) (*CarrierStore, error) { + if err := db.AutoMigrate(&CarrierRow{}); err != nil { + return nil, fmt.Errorf("migrating cluster_carrier: %w", err) + } + return &CarrierStore{db: db}, nil +} + +// Get returns the row, or ErrNotSeeded. +func (s *CarrierStore) Get(ctx context.Context) (CarrierRow, error) { + var row CarrierRow + err := s.db.WithContext(ctx).Where("id = ?", carrierRowID).Take(&row).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return CarrierRow{}, ErrNotSeeded + } + if err != nil { + return CarrierRow{}, fmt.Errorf("reading cluster carrier: %w", err) + } + return row, nil +} + +// Seed inserts a stable row at epoch 1 when none exists, and returns the row +// that is in the database after the call. created is true only for the caller +// whose insert won. Replicas that start together can all call it: the first +// insert wins and the others read that row. +func (s *CarrierStore) Seed(ctx context.Context, active Carrier, by string) (row CarrierRow, created bool, err error) { + if !active.valid() { + return CarrierRow{}, false, fmt.Errorf("%w: %q", ErrInvalidCarrier, active) + } + res := s.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(&CarrierRow{ + ID: carrierRowID, + Active: active, + Epoch: 1, + State: StateStable, + ChangedBy: by, + ChangedAt: time.Now().UTC(), + }) + if res.Error != nil { + return CarrierRow{}, false, fmt.Errorf("seeding cluster carrier: %w", res.Error) + } + row, err = s.Get(ctx) + return row, res.RowsAffected == 1, err +} + +// Transition writes change if the row is still at epoch from, and bumps the +// epoch. It returns ErrStaleEpoch when another writer got there first, so two +// admins or two replicas cannot interleave a change. +func (s *CarrierStore) Transition(ctx context.Context, from int64, change Change) (CarrierRow, error) { + if err := change.validate(); err != nil { + return CarrierRow{}, err + } + // RETURNING hands back the row this update wrote, not a later one. + var row CarrierRow + res := s.db.WithContext(ctx).Model(&row).Clauses(clause.Returning{}). + Where("id = ? AND epoch = ?", carrierRowID, from). + Updates(map[string]any{ + "active": change.Active, + "state": change.State, + "target": change.Target, + "draining": change.Draining, + "draining_until": change.DrainingUntil, + "changed_by": change.By, + "changed_at": time.Now().UTC(), + "prev_epoch": gorm.Expr("epoch"), + "epoch": gorm.Expr("epoch + 1"), + }) + if res.Error != nil { + return CarrierRow{}, fmt.Errorf("changing cluster carrier: %w", res.Error) + } + if res.RowsAffected == 0 { + // Tell a missing row from a lost race. + if _, err := s.Get(ctx); err != nil { + return CarrierRow{}, err + } + return CarrierRow{}, ErrStaleEpoch + } + return row, nil +} diff --git a/core/services/cluster/carrier_store_test.go b/core/services/cluster/carrier_store_test.go new file mode 100644 index 000000000000..79fe8edfd872 --- /dev/null +++ b/core/services/cluster/carrier_store_test.go @@ -0,0 +1,236 @@ +package cluster_test + +import ( + "context" + "sync" + "sync/atomic" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("CarrierStore", func() { + var ( + db *gorm.DB + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + }) + + newStore := func() *cluster.CarrierStore { + s, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + return s + } + + Describe("Get", func() { + It("reports a database that was never seeded", func() { + _, err := newStore().Get(ctx) + Expect(err).To(MatchError(cluster.ErrNotSeeded)) + }) + }) + + Describe("Seed", func() { + It("inserts a stable row at epoch 1", func() { + row, created, err := newStore().Seed(ctx, cluster.CarrierNATS, "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(created).To(BeTrue()) + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + Expect(row.Epoch).To(Equal(int64(1))) + Expect(row.State).To(Equal(cluster.StateStable)) + Expect(row.Target).To(BeEmpty()) + Expect(row.ChangedBy).To(Equal("replica-a")) + }) + + It("keeps the existing row when it is called again", func() { + _, _, err := newStore().Seed(ctx, cluster.CarrierNATS, "replica-a") + Expect(err).ToNot(HaveOccurred()) + + row, created, err := newStore().Seed(ctx, cluster.CarrierTunnel, "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(created).To(BeFalse()) + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + Expect(row.ChangedBy).To(Equal("replica-a")) + Expect(row.Epoch).To(Equal(int64(1))) + }) + + It("refuses a carrier name it does not know", func() { + _, _, err := newStore().Seed(ctx, cluster.Carrier("carrier-pigeon"), "replica-a") + Expect(err).To(MatchError(cluster.ErrInvalidCarrier)) + }) + + It("lets exactly one of several racing replicas insert the row", func() { + const replicas = 8 + stores := make([]*cluster.CarrierStore, replicas) + for i := range stores { + stores[i] = newStore() + } + + var created atomic.Int32 + rows := make([]cluster.CarrierRow, replicas) + errs := make([]error, replicas) + start := make(chan struct{}) + var wg sync.WaitGroup + for i := range stores { + wg.Add(1) + go func() { + defer wg.Done() + <-start + // Half the replicas have a NATS URL, half do not: the row + // must still come from one writer only. + want := cluster.CarrierNATS + if i%2 == 1 { + want = cluster.CarrierTunnel + } + var made bool + rows[i], made, errs[i] = stores[i].Seed(ctx, want, "replica") + if made { + created.Add(1) + } + }() + } + close(start) + wg.Wait() + + Expect(created.Load()).To(Equal(int32(1))) + for i := range rows { + Expect(errs[i]).ToNot(HaveOccurred()) + Expect(rows[i].Active).To(Equal(rows[0].Active), "every replica sees the same winner") + Expect(rows[i].Epoch).To(Equal(int64(1))) + } + var count int64 + Expect(db.Table("cluster_carrier").Count(&count).Error).To(Succeed()) + Expect(count).To(Equal(int64(1))) + }) + }) + + Describe("single row", func() { + It("refuses a second row at the database", func() { + _, _, err := newStore().Seed(ctx, cluster.CarrierNATS, "replica-a") + Expect(err).ToNot(HaveOccurred()) + err = db.Exec(`INSERT INTO cluster_carrier (id, active, epoch, state) VALUES (2, 'nats', 1, 'stable')`).Error + Expect(err).To(HaveOccurred()) + }) + }) + + Describe("Transition", func() { + var store *cluster.CarrierStore + + BeforeEach(func() { + store = newStore() + _, _, err := store.Seed(ctx, cluster.CarrierNATS, "replica-a") + Expect(err).ToNot(HaveOccurred()) + }) + + It("moves the row and bumps the epoch when the epoch matches", func() { + row, err := store.Transition(ctx, 1, cluster.Change{ + Active: cluster.CarrierNATS, + State: cluster.StatePrepare, + Target: cluster.CarrierTunnel, + By: "admin", + }) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Epoch).To(Equal(int64(2))) + Expect(row.PrevEpoch).To(Equal(int64(1))) + Expect(row.State).To(Equal(cluster.StatePrepare)) + Expect(row.Target).To(Equal(cluster.CarrierTunnel)) + Expect(row.ChangedBy).To(Equal("admin")) + + got, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Epoch).To(Equal(int64(2))) + Expect(got.Target).To(Equal(cluster.CarrierTunnel)) + }) + + It("refuses a stale epoch and leaves the row alone", func() { + _, err := store.Transition(ctx, 1, cluster.Change{ + Active: cluster.CarrierNATS, State: cluster.StatePrepare, Target: cluster.CarrierTunnel, By: "admin-1", + }) + Expect(err).ToNot(HaveOccurred()) + + _, err = store.Transition(ctx, 1, cluster.Change{ + Active: cluster.CarrierTunnel, State: cluster.StateStable, By: "admin-2", + }) + Expect(err).To(MatchError(cluster.ErrStaleEpoch)) + + got, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Epoch).To(Equal(int64(2))) + Expect(got.ChangedBy).To(Equal("admin-1")) + Expect(got.Active).To(Equal(cluster.CarrierNATS)) + }) + + It("lets one of two writers that read the same epoch win", func() { + other := newStore() + var wins, stale atomic.Int32 + start := make(chan struct{}) + var wg sync.WaitGroup + for _, s := range []*cluster.CarrierStore{store, other} { + wg.Add(1) + go func() { + defer wg.Done() + <-start + _, err := s.Transition(ctx, 1, cluster.Change{ + Active: cluster.CarrierNATS, State: cluster.StatePrepare, Target: cluster.CarrierTunnel, By: "admin", + }) + switch { + case err == nil: + wins.Add(1) + case err == cluster.ErrStaleEpoch: + stale.Add(1) + default: + Fail(err.Error()) + } + }() + } + close(start) + wg.Wait() + Expect(wins.Load()).To(Equal(int32(1))) + Expect(stale.Load()).To(Equal(int32(1))) + + got, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Epoch).To(Equal(int64(2))) + }) + + It("keeps the epoch strictly increasing over a chain of transitions", func() { + epoch := int64(1) + for range 5 { + row, err := store.Transition(ctx, epoch, cluster.Change{ + Active: cluster.CarrierNATS, State: cluster.StateStable, By: "admin", + }) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Epoch).To(Equal(epoch + 1)) + Expect(row.PrevEpoch).To(Equal(epoch)) + epoch = row.Epoch + } + }) + + It("refuses a change that does not describe a valid state", func() { + for _, bad := range []cluster.Change{ + {Active: cluster.Carrier("x"), State: cluster.StateStable}, + {Active: cluster.CarrierNATS, State: "sideways"}, + {Active: cluster.CarrierNATS, State: cluster.StatePrepare}, // no target + {Active: cluster.CarrierNATS, State: cluster.StateStable, Target: cluster.CarrierTunnel}, // target while stable + } { + _, err := store.Transition(ctx, 1, bad) + Expect(err).To(MatchError(cluster.ErrInvalidChange), "%+v", bad) + } + got, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Epoch).To(Equal(int64(1))) + }) + + It("reports a database that was never seeded", func() { + Expect(db.Exec(`DELETE FROM cluster_carrier`).Error).To(Succeed()) + _, err := store.Transition(ctx, 1, cluster.Change{Active: cluster.CarrierNATS, State: cluster.StateStable}) + Expect(err).To(MatchError(cluster.ErrNotSeeded)) + }) + }) +}) diff --git a/core/services/cluster/cluster_suite_test.go b/core/services/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..a81b9782485e --- /dev/null +++ b/core/services/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCluster(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster test suite") +} From 14e922ce6139d356856ef5431ba58aff2de00043 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:18:47 +0000 Subject: [PATCH 002/162] refactor(distributed): take NodeControl where a concrete sender was used The backend and model managers and the replica reconciler took the NATS sender by its concrete type. Add NodeControl, the interface of every control call they make, and take it instead. The sender implements it, so nothing changes at run time. NodeControl names the three unload interfaces that the model loader reads by type assertion. A sender that lacked one would still compile and would silently lose that behaviour. Export the install-with-force fallback so a sender in another package can implement it. Add a ClientFactory option to the reconciler, so its default probe dials the way the router does. Add NewTokenClientFactory for the direct factory. Narrow the agent event bridge from MessagingClient to Broadcaster. It only publishes and subscribes. Pin that the NATS client returns a client for a server that is not up yet, which startup relies on. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/agents/events.go | 4 +-- .../services/messaging/client_connect_test.go | 25 ++++++++++++++ core/services/nodes/client_factory_test.go | 13 ++++++++ core/services/nodes/interfaces.go | 33 +++++++++++++++++++ core/services/nodes/managers_distributed.go | 10 +++--- core/services/nodes/reconciler.go | 22 +++++++++---- core/services/nodes/unloader.go | 5 +-- 7 files changed, 96 insertions(+), 16 deletions(-) create mode 100644 core/services/messaging/client_connect_test.go diff --git a/core/services/agents/events.go b/core/services/agents/events.go index b187b295ad25..57792ef137f8 100644 --- a/core/services/agents/events.go +++ b/core/services/agents/events.go @@ -40,7 +40,7 @@ type AgentCancelEvent struct { // EventBridge bridges agent events between NATS and SSE connections. // It enables cross-instance SSE: user connects to Frontend 1, agent runs on Frontend 2. type EventBridge struct { - nats messaging.MessagingClient + nats messaging.Broadcaster store *AgentStore instanceID string @@ -52,7 +52,7 @@ type EventBridge struct { } // NewEventBridge creates a new EventBridge. -func NewEventBridge(nc messaging.MessagingClient, store *AgentStore, instanceID string) *EventBridge { +func NewEventBridge(nc messaging.Broadcaster, store *AgentStore, instanceID string) *EventBridge { return &EventBridge{ nats: nc, store: store, diff --git a/core/services/messaging/client_connect_test.go b/core/services/messaging/client_connect_test.go new file mode 100644 index 000000000000..d1267a01bd09 --- /dev/null +++ b/core/services/messaging/client_connect_test.go @@ -0,0 +1,25 @@ +package messaging_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// A frontend that starts before its NATS server must come up and connect when +// the server arrives. Startup relies on this: New returns a client for a URL +// nothing listens on, and refuses only a URL it cannot parse. +var _ = Describe("Client connect", func() { + It("returns a client that keeps retrying when the server is not reachable", func() { + c, err := messaging.New("nats://127.0.0.1:1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(c.Close) + Expect(c.IsConnected()).To(BeFalse()) + }) + + It("refuses a URL it cannot parse", func() { + _, err := messaging.New("nats://[::1") + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/nodes/client_factory_test.go b/core/services/nodes/client_factory_test.go index 3c8dff7afc72..ec7d115daef2 100644 --- a/core/services/nodes/client_factory_test.go +++ b/core/services/nodes/client_factory_test.go @@ -78,4 +78,17 @@ var _ = Describe("Backend client construction carries the node id", func() { Expect(f.calls()).NotTo(BeEmpty()) Expect(f.calls()).To(HaveEach("n1@10.0.0.2:9001")) }) + + It("hands the node id and the replica address to the factory from the reconciler's probe", func() { + f := &recordingFactory{next: func() grpc.Backend { return &fakeBackendClient{healthy: true} }} + rc := NewReplicaReconciler(ReplicaReconcilerOptions{ClientFactory: f}) + + Expect(rc.prober.Probe(context.Background(), "n1", "10.0.0.3:9001")).To(Equal(ProbeAlive)) + Expect(f.calls()).To(Equal([]string{"n1@10.0.0.3:9001"})) + }) + + It("probes with a direct client when the reconciler is given no factory", func() { + rc := NewReplicaReconciler(ReplicaReconcilerOptions{RegistrationToken: "t"}) + Expect(rc.prober).To(Equal(grpcModelProber{clients: NewTokenClientFactory("t")})) + }) }) diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 52ca4321515d..5719dd48834b 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -7,6 +7,7 @@ import ( "github.com/mudler/LocalAI/core/services/workerctl" grpc "github.com/mudler/LocalAI/pkg/grpc" + "github.com/mudler/LocalAI/pkg/model" ) type ExactModelStopper interface { @@ -157,6 +158,31 @@ type LoadOperationControl interface { ExactModelStopper } +// NodeControl is every control verb the frontend sends to workers, whatever +// carries it. The router, the reconciler, the backend and model managers and +// the model loader take it instead of a concrete sender, so a carrier is +// replaced in one place. A carrier implements all of it; the part beyond +// NodeCommandSender is the lookups and fan-outs built on those verbs. +// +// The model loader reads the three unload interfaces through type assertions +// (see pkg/model), which is why each one is named here: a value that lacks one +// compiles and silently loses that behaviour. +type NodeControl interface { + NodeCommandSender + NodeProcessLister + model.RemoteModelUnloader + model.RemoteModelContextUnloader + model.RemoteModelPresenceChecker + + // InstallTimeout is the longest an install request waits for its reply. + InstallTimeout() time.Duration + // DeleteModelFiles asks every node that holds the model to delete its files. + DeleteModelFiles(modelName string) error + // InstallBackendForce is the rolling-update fallback for UpgradeBackend on + // a worker that does not serve backend upgrades: an install with Force set. + InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) +} + // ConcurrencyConflictResolver returns the names of configured models that // share at least one concurrency group with the given model. It is satisfied // by *config.ModelConfigLoader and lets the SmartRouter make group-aware @@ -247,6 +273,13 @@ func (f *tokenClientFactory) NewClient(_, address string, parallel bool) grpc.Ba return grpc.NewClient(address, parallel, nil, false) } +// NewTokenClientFactory returns the direct BackendClientFactory: it dials the +// address it is handed and sends token as a bearer credential when it is not +// empty. +func NewTokenClientFactory(token string) BackendClientFactory { + return &tokenClientFactory{token: token} +} + // WorkerNetDialerFor returns the dial function that reaches one worker's own // HTTP server, in the shape http.Transport.DialContext and // websocket.Dialer.NetDialContext take. It is keyed by node id, not address, diff --git a/core/services/nodes/managers_distributed.go b/core/services/nodes/managers_distributed.go index 0ea607cf0a46..a44629a7115a 100644 --- a/core/services/nodes/managers_distributed.go +++ b/core/services/nodes/managers_distributed.go @@ -20,13 +20,13 @@ import ( // for model deletion so worker nodes clean up stale files. type DistributedModelManager struct { local galleryop.ModelManager - adapter *RemoteUnloaderAdapter + adapter NodeControl } // NewDistributedModelManager creates a DistributedModelManager. // Backend auto-install is disabled because the frontend node delegates // inference to workers and never runs backends locally. -func NewDistributedModelManager(appConfig *config.ApplicationConfig, ml *model.ModelLoader, adapter *RemoteUnloaderAdapter) *DistributedModelManager { +func NewDistributedModelManager(appConfig *config.ApplicationConfig, ml *model.ModelLoader, adapter NodeControl) *DistributedModelManager { local := galleryop.NewLocalModelManager(appConfig, ml) local.SetAutoInstallBackend(false) return &DistributedModelManager{ @@ -59,7 +59,7 @@ type nodeProgressSink interface { // for backend deletion so worker nodes clean up stale files. type DistributedBackendManager struct { local galleryop.BackendManager - adapter *RemoteUnloaderAdapter + adapter NodeControl registry *NodeRegistry backendGalleries []config.Gallery systemState *system.SystemState @@ -69,7 +69,7 @@ type DistributedBackendManager struct { // NewDistributedBackendManager creates a DistributedBackendManager. // progressSink may be nil to disable per-node OpStatus writes (single-node // tests don't need it). -func NewDistributedBackendManager(appConfig *config.ApplicationConfig, ml *model.ModelLoader, adapter *RemoteUnloaderAdapter, registry *NodeRegistry, progressSink nodeProgressSink) *DistributedBackendManager { +func NewDistributedBackendManager(appConfig *config.ApplicationConfig, ml *model.ModelLoader, adapter NodeControl, registry *NodeRegistry, progressSink nodeProgressSink) *DistributedBackendManager { return &DistributedBackendManager{ local: galleryop.NewLocalBackendManager(appConfig, ml), adapter: adapter, @@ -601,7 +601,7 @@ func (d *DistributedBackendManager) UpgradeBackend(ctx context.Context, op *gall // Rolling-update fallback: an older worker doesn't know // backend.upgrade. Try the legacy install-with-force path. if errors.Is(err, ErrNoRoute) { - instReply, instErr := d.adapter.installWithForceFallback(node.ID, name, string(galleriesJSON), "", "", "", 0, opID, onProgressArg) + instReply, instErr := d.adapter.InstallBackendForce(node.ID, name, string(galleriesJSON), "", "", "", 0, opID, onProgressArg) if instErr != nil { return instErr } diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index 214245351c5b..7b2a2d9ccb42 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -43,7 +43,7 @@ type ModelProber interface { } // NodeProcessLister asks a worker which model backend processes it currently -// has running. Implemented by RemoteUnloaderAdapter over NATS. +// has running. NodeControl carries it. // // This is the sounder liveness signal: the worker owns the process table, so // its answer does not depend on whether a backend is busy. A health probe @@ -121,8 +121,8 @@ type ReplicaReconciler struct { registry *NodeRegistry scheduler ModelScheduler // interface for scheduling new models unloader NodeCommandSender - adapter *RemoteUnloaderAdapter // NATS sender for pending-op drain - prober ModelProber // health probe for model gRPC addrs + adapter NodeControl // sender for pending-op drain + prober ModelProber // health probe for model gRPC addrs db *gorm.DB interval time.Duration scaleDownDelay time.Duration @@ -172,14 +172,18 @@ type ReplicaReconcilerOptions struct { Registry *NodeRegistry Scheduler ModelScheduler Unloader NodeCommandSender - // Adapter is the NATS sender used to retry pending backend ops. When nil, + // Adapter is the control carrier used to retry pending backend ops. When nil, // the state-reconciler pending-drain pass is a no-op (single-node mode). - Adapter *RemoteUnloaderAdapter + Adapter NodeControl // RegistrationToken is used by the default gRPC prober when probing model // addresses. Matches the worker's token so HealthCheck auth succeeds. RegistrationToken string // Prober overrides the default gRPC health probe (used by tests). Prober ModelProber + // ClientFactory builds the clients of the default prober, so probes dial + // the way the router does. When nil it dials the address directly with + // RegistrationToken. + ClientFactory BackendClientFactory // ProcessLister overrides the default worker process query. When nil and // no Adapter is set, the worker-authoritative pass is skipped entirely and // only the port probe runs. @@ -218,7 +222,11 @@ func NewReplicaReconciler(opts ReplicaReconcilerOptions) *ReplicaReconciler { } prober := opts.Prober if prober == nil { - prober = grpcModelProber{clients: &tokenClientFactory{token: opts.RegistrationToken}} + clients := opts.ClientFactory + if clients == nil { + clients = &tokenClientFactory{token: opts.RegistrationToken} + } + prober = grpcModelProber{clients: clients} } pressureThreshold := opts.PressureThreshold if pressureThreshold == 0 { @@ -358,7 +366,7 @@ func (rc *ReplicaReconciler) drainPendingBackendOps(ctx context.Context) { reply, err := rc.adapter.UpgradeBackend(op.NodeID, op.Backend, string(op.Galleries), "", "", "", 0, "", nil) if err != nil { if errors.Is(err, ErrNoRoute) { - instReply, instErr := rc.adapter.installWithForceFallback(op.NodeID, op.Backend, string(op.Galleries), "", "", "", 0, "", nil) + instReply, instErr := rc.adapter.InstallBackendForce(op.NodeID, op.Backend, string(op.Galleries), "", "", "", 0, "", nil) if instErr != nil { applyErr = instErr } else if !instReply.Success { diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 6682437a7e38..53ef79834dc1 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -89,6 +89,7 @@ var ( _ model.RemoteModelContextUnloader = (*RemoteUnloaderAdapter)(nil) _ model.RemoteModelPresenceChecker = (*RemoteUnloaderAdapter)(nil) _ ExactModelStopper = (*RemoteUnloaderAdapter)(nil) + _ NodeControl = (*RemoteUnloaderAdapter)(nil) ) const exactModelStopTimeout = 10 * time.Second @@ -373,13 +374,13 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO return reply, err } -// installWithForceFallback is the rolling-update fallback used by +// InstallBackendForce is the rolling-update fallback used by // DistributedBackendManager.UpgradeBackend when backend.upgrade returns // ErrNoRoute (the worker is on a pre-2026-05-08 build that // doesn't subscribe to the new subject). It re-fires the legacy // backend.install with Force=true. Drop this once every worker is on // 2026-05-08 or newer. -func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (a *RemoteUnloaderAdapter) InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { subject := messaging.SubjectNodeBackendInstall(nodeID) xlog.Warn("Falling back to legacy backend.install Force=true (old worker)", "nodeID", nodeID, "backend", backendType) From d4418592be819ec11e166de347f53f750ef59313 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:18:47 +0000 Subject: [PATCH 003/162] feat(distributed): add forwarding holders for the carrier seams A carrier.Set is the implementation of every seam for one carrier. The holders implement the same interfaces as the seams (fan-out, work queue, control verbs, file staging, backend clients, worker dialer, agent RPC) and forward each call to the set that an atomic pointer names. They take no lock and add no allocation. The fan-out holder keeps a registry of its subscriptions. Listen attaches every subscription to a second set, so each handler listens on both carriers while a swap runs. Release drops the first set again. Publish always uses the set the pointer names. The holder also forwards the reconnect hook that syncstate and the failover sync register. NewNATSSet builds the set for the NATS carrier. A spec fails when any other non-test file imports the NATS libraries. Nothing stores a second set yet, so no behaviour changes. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/broadcaster.go | 238 +++++++++++ core/services/carrier/broadcaster_test.go | 395 ++++++++++++++++++ core/services/carrier/carrier_suite_test.go | 13 + core/services/carrier/fakes_test.go | 230 ++++++++++ core/services/carrier/holders.go | 192 +++++++++ core/services/carrier/holders_test.go | 310 ++++++++++++++ core/services/carrier/import_boundary_test.go | 184 ++++++++ core/services/carrier/nats.go | 78 ++++ core/services/carrier/nats_test.go | 95 +++++ core/services/carrier/set.go | 81 ++++ 10 files changed, 1816 insertions(+) create mode 100644 core/services/carrier/broadcaster.go create mode 100644 core/services/carrier/broadcaster_test.go create mode 100644 core/services/carrier/carrier_suite_test.go create mode 100644 core/services/carrier/fakes_test.go create mode 100644 core/services/carrier/holders.go create mode 100644 core/services/carrier/holders_test.go create mode 100644 core/services/carrier/import_boundary_test.go create mode 100644 core/services/carrier/nats.go create mode 100644 core/services/carrier/nats_test.go create mode 100644 core/services/carrier/set.go diff --git a/core/services/carrier/broadcaster.go b/core/services/carrier/broadcaster.go new file mode 100644 index 000000000000..882059daf9a5 --- /dev/null +++ b/core/services/carrier/broadcaster.go @@ -0,0 +1,238 @@ +package carrier + +import ( + "errors" + "fmt" + "slices" + "sync" + "sync/atomic" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// ErrActiveSet is returned when a caller asks to release the active set. +var ErrActiveSet = errors.New("the active carrier set cannot be released") + +// Broadcaster is the fan-out holder. Publish goes to the current set. It also +// keeps a registry of every subscription, because a subscription lives on one +// carrier and has to be made again on the next one. +// +// A swap runs in three calls: +// +// - Listen(next) attaches every registered subscription to next as well, so +// each handler listens on both carriers. Publish still uses the old one. +// - The caller stores next in the pointer. Publish now uses next, and a +// message that a slower peer still publishes on the old carrier is heard. +// - Release(old) drops the subscriptions on the old carrier. +// +// No message is published on two carriers, so nothing is delivered twice. +type Broadcaster struct { + cur *atomic.Pointer[Set] + + mu sync.Mutex + listening []*Set // sets the subscriptions are attached to + hooked map[*Set]struct{} + subs map[*subscription]struct{} + reconnect []func() +} + +var _ messaging.Broadcaster = (*Broadcaster)(nil) + +// NewBroadcaster returns a holder over cur, which must name a set. Existing +// and later subscriptions attach to that set. +func NewBroadcaster(cur *atomic.Pointer[Set]) *Broadcaster { + b := &Broadcaster{ + cur: cur, + hooked: map[*Set]struct{}{}, + subs: map[*subscription]struct{}{}, + } + first := cur.Load() + b.listening = []*Set{first} + b.hook(first) + return b +} + +// Publish sends on the current set. It takes no lock. +func (b *Broadcaster) Publish(subject string, data any) error { + return b.cur.Load().Broadcaster.Publish(subject, data) +} + +// Subscribe attaches handler to every set the holder listens on. It fails, and +// leaves nothing attached, when any of them refuses. +func (b *Broadcaster) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { + sub := &subscription{owner: b, subject: subject, handler: handler, inner: map[*Set]messaging.Subscription{}} + + // Read the sets and register in one step, so a concurrent Listen either + // sees this subscription or is seen by it, never neither. + b.mu.Lock() + sets := slices.Clone(b.listening) + b.subs[sub] = struct{}{} + b.mu.Unlock() + + for _, s := range sets { + if err := sub.attach(s); err != nil { + _ = sub.Unsubscribe() + return nil, err + } + } + return sub, nil +} + +// Listen attaches every registered subscription to next. It does nothing if +// next is already attached. On failure it detaches next again, so the caller +// can report that the carrier is not ready. +func (b *Broadcaster) Listen(next *Set) error { + b.mu.Lock() + if slices.Contains(b.listening, next) { + b.mu.Unlock() + return nil + } + b.listening = append(b.listening, next) + b.hook(next) + subs := b.snapshotLocked() + b.mu.Unlock() + + for _, sub := range subs { + if err := sub.attach(next); err != nil { + _ = b.Release(next) + return fmt.Errorf("attaching subscriptions to carrier %q: %w", next.Name, err) + } + } + return nil +} + +// Release detaches every subscription from old. It refuses the active set. +// Releasing a set that is not attached does nothing. +func (b *Broadcaster) Release(old *Set) error { + if b.cur.Load() == old { + return ErrActiveSet + } + b.mu.Lock() + i := slices.Index(b.listening, old) + if i < 0 { + b.mu.Unlock() + return nil + } + b.listening = slices.Delete(b.listening, i, i+1) + subs := b.snapshotLocked() + b.mu.Unlock() + + var errs error + for _, sub := range subs { + errs = errors.Join(errs, sub.detach(old)) + } + return errs +} + +// OnReconnect registers cb to run after any carrier the holder uses recovers a +// lost connection. Consumers re-read their state from tables on it, because a +// reconnecting subscriber misses messages. A nil cb is ignored. +func (b *Broadcaster) OnReconnect(cb func()) { + if cb == nil { + return + } + b.mu.Lock() + b.reconnect = append(b.reconnect, cb) + b.mu.Unlock() +} + +// NotifyReconnect runs every registered hook once. The swap code calls it when +// a swap is complete: messages sent across the flip are not ordered, so the +// consumers read their state again. +func (b *Broadcaster) NotifyReconnect() { + b.mu.Lock() + cbs := slices.Clone(b.reconnect) + b.mu.Unlock() + for _, cb := range cbs { + cb() + } +} + +// hook asks a set to tell the holder about its reconnects. The caller holds +// b.mu or owns b exclusively. A set is hooked once for the life of the holder. +func (b *Broadcaster) hook(s *Set) { + if _, done := b.hooked[s]; done || s.OnReconnect == nil { + return + } + b.hooked[s] = struct{}{} + s.OnReconnect(b.NotifyReconnect) +} + +func (b *Broadcaster) snapshotLocked() []*subscription { + subs := make([]*subscription, 0, len(b.subs)) + for s := range b.subs { + subs = append(subs, s) + } + return subs +} + +func (b *Broadcaster) attached(s *Set) bool { + b.mu.Lock() + defer b.mu.Unlock() + return slices.Contains(b.listening, s) +} + +// subscription is one registered subscription and its attachment to each set. +type subscription struct { + owner *Broadcaster + subject string + handler func([]byte) + + mu sync.Mutex + closed bool + inner map[*Set]messaging.Subscription +} + +// attach subscribes on s. It drops the new subscription again when this one +// was closed, or s was released, while the carrier was busy. +func (sub *subscription) attach(s *Set) error { + in, err := s.Broadcaster.Subscribe(sub.subject, sub.handler) + if err != nil { + return err + } + sub.mu.Lock() + defer sub.mu.Unlock() + if sub.closed || !sub.owner.attached(s) { + return in.Unsubscribe() + } + if _, dup := sub.inner[s]; dup { + return in.Unsubscribe() + } + sub.inner[s] = in + return nil +} + +func (sub *subscription) detach(s *Set) error { + sub.mu.Lock() + in, ok := sub.inner[s] + delete(sub.inner, s) + sub.mu.Unlock() + if !ok { + return nil + } + return in.Unsubscribe() +} + +// Unsubscribe removes the subscription from every set. It is safe to call +// twice. +func (sub *subscription) Unsubscribe() error { + sub.mu.Lock() + if sub.closed { + sub.mu.Unlock() + return nil + } + sub.closed = true + inner := sub.inner + sub.inner = nil + sub.mu.Unlock() + + sub.owner.mu.Lock() + delete(sub.owner.subs, sub) + sub.owner.mu.Unlock() + + var errs error + for _, in := range inner { + errs = errors.Join(errs, in.Unsubscribe()) + } + return errs +} diff --git a/core/services/carrier/broadcaster_test.go b/core/services/carrier/broadcaster_test.go new file mode 100644 index 000000000000..a1160bb42079 --- /dev/null +++ b/core/services/carrier/broadcaster_test.go @@ -0,0 +1,395 @@ +package carrier_test + +import ( + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// received collects the payloads a handler saw. +type received struct { + mu sync.Mutex + data []string +} + +func (r *received) handler() func([]byte) { + return func(b []byte) { + r.mu.Lock() + r.data = append(r.data, string(b)) + r.mu.Unlock() + } +} + +func (r *received) all() []string { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.data...) +} + +var _ = Describe("Broadcaster holder", func() { + var ( + a, b *fakeCarrier + cur atomic.Pointer[carrier.Set] + h *carrier.Broadcaster + ) + + BeforeEach(func() { + a = newFakeCarrier(cluster.CarrierNATS, 1) + b = newFakeCarrier(cluster.CarrierTunnel, 2) + cur = atomic.Pointer[carrier.Set]{} + cur.Store(a.set) + h = carrier.NewBroadcaster(&cur) + }) + + It("publishes on the carrier the pointer names", func() { + got := &received{} + _, err := a.bus.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + _, err = b.bus.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + + Expect(h.Publish("nodes.x", "one")).To(Succeed()) + Expect(got.all()).To(HaveLen(1)) + Expect(a.bus.PublishCount("nodes.x")).To(Equal(1)) + Expect(b.bus.PublishCount("nodes.x")).To(Equal(0)) + + cur.Store(b.set) + Expect(h.Publish("nodes.x", "two")).To(Succeed()) + Expect(a.bus.PublishCount("nodes.x")).To(Equal(1)) + Expect(b.bus.PublishCount("nodes.x")).To(Equal(1)) + }) + + It("delivers a subscription that was made before and after a swap", func() { + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + + Expect(h.Listen(b.set)).To(Succeed()) + cur.Store(b.set) + + // A peer that has not flipped yet still publishes on the old carrier. + Expect(a.bus.Publish("nodes.x", "from-old")).To(Succeed()) + Expect(h.Publish("nodes.x", "from-new")).To(Succeed()) + + Expect(got.all()).To(ConsistOf(`"from-old"`, `"from-new"`)) + }) + + It("listens on both carriers while a target is attached, and publishes on one", func() { + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + Expect(h.Listen(b.set)).To(Succeed()) + + Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) + Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(got.all()).To(ConsistOf(`"a"`, `"b"`)) + + Expect(h.Publish("nodes.x", "c")).To(Succeed()) + Expect(a.bus.PublishCount("nodes.x")).To(Equal(2), "the active carrier carries the publish") + Expect(b.bus.PublishCount("nodes.x")).To(Equal(1), "the target carrier carries none") + }) + + It("subscribes on every attached carrier when a subscription is made during the window", func() { + Expect(h.Listen(b.set)).To(Succeed()) + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + + Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) + Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(got.all()).To(ConsistOf(`"a"`, `"b"`)) + }) + + It("removes a subscription from every carrier when it is unsubscribed", func() { + got := &received{} + sub, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + Expect(h.Listen(b.set)).To(Succeed()) + + Expect(sub.Unsubscribe()).To(Succeed()) + Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) + Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(got.all()).To(BeEmpty()) + + // It stays gone for a carrier attached later. + c := newFakeCarrier(cluster.CarrierTunnel, 3) + Expect(h.Listen(c.set)).To(Succeed()) + Expect(c.bus.Publish("nodes.x", "c")).To(Succeed()) + Expect(got.all()).To(BeEmpty()) + }) + + It("tolerates a second Unsubscribe", func() { + sub, err := h.Subscribe("nodes.x", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(sub.Unsubscribe()).To(Succeed()) + Expect(sub.Unsubscribe()).To(Succeed()) + }) + + It("drops the old carrier's subscriptions on Release and keeps the new ones", func() { + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + Expect(h.Listen(b.set)).To(Succeed()) + cur.Store(b.set) + + Expect(h.Release(a.set)).To(Succeed()) + Expect(a.bus.Publish("nodes.x", "late-old")).To(Succeed()) + Expect(b.bus.Publish("nodes.x", "new")).To(Succeed()) + Expect(got.all()).To(Equal([]string{`"new"`})) + }) + + It("refuses to release the carrier that is active", func() { + Expect(h.Release(a.set)).To(MatchError(carrier.ErrActiveSet)) + }) + + It("treats a second Listen of the same carrier as nothing to do", func() { + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + Expect(h.Listen(b.set)).To(Succeed()) + Expect(h.Listen(b.set)).To(Succeed()) + + Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(got.all()).To(HaveLen(1), "one subscription per carrier, not two") + }) + + It("leaves nothing behind when a carrier cannot take a subscription", func() { + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).ToNot(HaveOccurred()) + + broken := newFakeCarrier(cluster.CarrierTunnel, 3) + boom := errors.New("subscribe refused") + broken.set.Broadcaster = failingBroadcaster{Broadcaster: broken.bus, err: boom} + + err = h.Listen(broken.set) + Expect(err).To(MatchError(boom)) + + // The target is not attached, and a later subscription does not try it. + _, err = h.Subscribe("nodes.y", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(h.Release(broken.set)).To(Succeed(), "releasing an unattached carrier is harmless") + }) + + It("fails a subscription when an attached carrier refuses it, and undoes the part that worked", func() { + refusing := newFakeCarrier(cluster.CarrierTunnel, 3) + boom := errors.New("subscribe refused") + refusing.set.Broadcaster = failingBroadcaster{Broadcaster: refusing.bus, err: boom} + Expect(h.Listen(refusing.set)).To(Succeed(), "nothing is registered yet, so nothing is refused") + + got := &received{} + _, err := h.Subscribe("nodes.x", got.handler()) + Expect(err).To(MatchError(boom)) + + Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) + Expect(got.all()).To(BeEmpty(), "the half-made subscription was removed") + }) + + Describe("reconnect hooks", func() { + It("runs a hook when the active carrier reconnects", func() { + var n atomic.Int32 + h.OnReconnect(func() { n.Add(1) }) + a.bus.TriggerReconnect() + Expect(n.Load()).To(Equal(int32(1))) + }) + + It("follows the hook to a carrier attached later", func() { + var n atomic.Int32 + h.OnReconnect(func() { n.Add(1) }) + Expect(h.Listen(b.set)).To(Succeed()) + cur.Store(b.set) + b.bus.TriggerReconnect() + Expect(n.Load()).To(Equal(int32(1))) + }) + + It("runs a hook once per reconnect however often a carrier is attached", func() { + var n atomic.Int32 + h.OnReconnect(func() { n.Add(1) }) + Expect(h.Listen(b.set)).To(Succeed()) + Expect(h.Release(b.set)).To(Succeed()) + Expect(h.Listen(b.set)).To(Succeed()) + b.bus.TriggerReconnect() + Expect(n.Load()).To(Equal(int32(1))) + }) + + It("ignores a nil hook", func() { + h.OnReconnect(nil) + a.bus.TriggerReconnect() + }) + + It("lets the swap code run every hook once", func() { + var n atomic.Int32 + h.OnReconnect(func() { n.Add(1) }) + h.OnReconnect(func() { n.Add(1) }) + h.NotifyReconnect() + Expect(n.Load()).To(Equal(int32(2))) + }) + + It("works with a carrier that has no reconnect notion", func() { + c := newFakeCarrier(cluster.CarrierTunnel, 3) + c.set.OnReconnect = nil + Expect(h.Listen(c.set)).To(Succeed()) + }) + }) + + Describe("while a swap is in progress", func() { + It("does not hold a publish behind a listen that is waiting on a carrier", func() { + slow := newFakeCarrier(cluster.CarrierTunnel, 3) + gate := make(chan struct{}) + entered := make(chan struct{}, 1) + slow.set.Broadcaster = blockingBroadcaster{Broadcaster: slow.bus, gate: gate, entered: entered} + + _, err := h.Subscribe("nodes.x", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + + listenDone := make(chan error, 1) + go func() { listenDone <- h.Listen(slow.set) }() + Eventually(entered).Should(Receive()) + + published := make(chan error, 1) + go func() { published <- h.Publish("nodes.x", "during") }() + Eventually(published, time.Second).Should(Receive(BeNil())) + + subscribed := make(chan error, 1) + go func() { + _, err := h.Subscribe("nodes.z", func([]byte) {}) + subscribed <- err + }() + Consistently(listenDone, 50*time.Millisecond).ShouldNot(Receive()) + + close(gate) + Eventually(listenDone, time.Second).Should(Receive(BeNil())) + Eventually(subscribed, time.Second).Should(Receive(BeNil())) + }) + + It("does not block a publish already running on the old carrier when the pointer moves", func() { + stuck := newFakeCarrier(cluster.CarrierNATS, 1) + gate := make(chan struct{}) + entered := make(chan struct{}, 1) + stuck.set.Broadcaster = blockingPublisher{Broadcaster: stuck.bus, gate: gate, entered: entered} + cur.Store(stuck.set) + + first := make(chan error, 1) + go func() { first <- h.Publish("nodes.x", "old") }() + Eventually(entered).Should(Receive()) + + cur.Store(b.set) + Expect(h.Publish("nodes.x", "new")).To(Succeed()) + Expect(b.bus.PublishCount("nodes.x")).To(Equal(1)) + + close(gate) + Eventually(first, time.Second).Should(Receive(BeNil())) + }) + + It("neither loses nor repeats a subscription that races a Listen", func() { + const n = 40 + var wg sync.WaitGroup + counts := make([]*received, n) + for i := range counts { + counts[i] = &received{} + } + start := make(chan struct{}) + for i := range n { + wg.Add(1) + go func() { + defer wg.Done() + <-start + _, err := h.Subscribe("nodes.x", counts[i].handler()) + Expect(err).ToNot(HaveOccurred()) + }() + } + wg.Add(1) + go func() { + defer wg.Done() + <-start + Expect(h.Listen(b.set)).To(Succeed()) + }() + close(start) + wg.Wait() + + Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) + Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + for i := range counts { + Expect(counts[i].all()).To(ConsistOf(`"a"`, `"b"`), "subscription %d", i) + } + }) + }) + + Describe("hot path", func() { + It("adds no allocation to a publish", func() { + quiet := &quietBroadcaster{} + c := newFakeCarrier(cluster.CarrierNATS, 1) + c.set.Broadcaster = quiet + var p atomic.Pointer[carrier.Set] + p.Store(c.set) + holder := carrier.NewBroadcaster(&p) + + var payload any = "payload" + allocs := testing.AllocsPerRun(1000, func() { + _ = holder.Publish("nodes.x", payload) + }) + Expect(allocs).To(BeZero()) + Expect(quiet.published.Load()).To(BeNumerically(">", 1000)) + }) + }) +}) + +var _ messaging.Broadcaster = failingBroadcaster{} + +// failingBroadcaster refuses Subscribe. +type failingBroadcaster struct { + messaging.Broadcaster + err error +} + +func (f failingBroadcaster) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { + return nil, f.err +} + +// blockingBroadcaster blocks Subscribe until gate closes. +type blockingBroadcaster struct { + messaging.Broadcaster + gate chan struct{} + entered chan struct{} +} + +func (f blockingBroadcaster) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { + select { + case f.entered <- struct{}{}: + default: + } + <-f.gate + return f.Broadcaster.Subscribe(subject, handler) +} + +// blockingPublisher blocks Publish until gate closes. +type blockingPublisher struct { + messaging.Broadcaster + gate chan struct{} + entered chan struct{} +} + +func (f blockingPublisher) Publish(subject string, data any) error { + select { + case f.entered <- struct{}{}: + default: + } + <-f.gate + return f.Broadcaster.Publish(subject, data) +} + +// quietBroadcaster counts publishes without allocating. +type quietBroadcaster struct{ published atomic.Int64 } + +func (q *quietBroadcaster) Publish(string, any) error { q.published.Add(1); return nil } +func (q *quietBroadcaster) Subscribe(string, func([]byte)) (messaging.Subscription, error) { + return nil, nil +} diff --git a/core/services/carrier/carrier_suite_test.go b/core/services/carrier/carrier_suite_test.go new file mode 100644 index 000000000000..089a3a10a121 --- /dev/null +++ b/core/services/carrier/carrier_suite_test.go @@ -0,0 +1,13 @@ +package carrier_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCarrier(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Carrier holders test suite") +} diff --git a/core/services/carrier/fakes_test.go b/core/services/carrier/fakes_test.go new file mode 100644 index 000000000000..06799a4dcdd7 --- /dev/null +++ b/core/services/carrier/fakes_test.go @@ -0,0 +1,230 @@ +package carrier_test + +import ( + "context" + "net" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +// recorder remembers which method a fake carrier member received. +type recorder struct { + mu sync.Mutex + calls []string +} + +func (r *recorder) rec(name string) { + r.mu.Lock() + r.calls = append(r.calls, name) + r.mu.Unlock() +} + +func (r *recorder) seen() []string { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.calls...) +} + +// fakeCommands is a nodes.NodeControl that only records the call. +type fakeCommands struct{ recorder } + +var _ nodes.NodeControl = (*fakeCommands)(nil) + +func (f *fakeCommands) InstallBackend(string, string, string, string, string, string, string, int, string, func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + f.rec("InstallBackend") + return &workerctl.BackendInstallReply{}, nil +} + +func (f *fakeCommands) UpgradeBackend(string, string, string, string, string, string, int, string, func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error) { + f.rec("UpgradeBackend") + return &workerctl.BackendUpgradeReply{}, nil +} + +func (f *fakeCommands) DeleteBackend(string, string) (*workerctl.BackendDeleteReply, error) { + f.rec("DeleteBackend") + return &workerctl.BackendDeleteReply{}, nil +} + +func (f *fakeCommands) ListBackends(string) (*workerctl.BackendListReply, error) { + f.rec("ListBackends") + return &workerctl.BackendListReply{}, nil +} +func (f *fakeCommands) StopBackend(string, string) error { f.rec("StopBackend"); return nil } +func (f *fakeCommands) UnloadModelOnNode(string, string) error { + f.rec("UnloadModelOnNode") + return nil +} +func (f *fakeCommands) PingNode(string) error { f.rec("PingNode"); return nil } +func (f *fakeCommands) InstallBackendOp(string, string, string, string, int, string, string, time.Duration, func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + f.rec("InstallBackendOp") + return &workerctl.BackendInstallReply{}, nil +} + +func (f *fakeCommands) StopLoadOperation(context.Context, string, workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { + f.rec("StopLoadOperation") + return workerctl.ModelStopReply{}, nil +} + +func (f *fakeCommands) OperationControl(string, workerctl.OperationRequest) (*workerctl.OperationReply, error) { + f.rec("OperationControl") + return &workerctl.OperationReply{}, nil +} +func (f *fakeCommands) UnloadReplica(string, nodes.NodeModel) error { + f.rec("UnloadReplica") + return nil +} + +func (f *fakeCommands) StopModelReplica(context.Context, string, nodes.NodeModel, bool) (workerctl.ModelStopReply, error) { + f.rec("StopModelReplica") + return workerctl.ModelStopReply{}, nil +} + +func (f *fakeCommands) ListRunningModels(string) (*workerctl.ModelsRunningReply, error) { + f.rec("ListRunningModels") + return &workerctl.ModelsRunningReply{}, nil +} +func (f *fakeCommands) UnloadRemoteModel(string) error { f.rec("UnloadRemoteModel"); return nil } +func (f *fakeCommands) UnloadRemoteModelContext(context.Context, string, bool) error { + f.rec("UnloadRemoteModelContext") + return nil +} + +func (f *fakeCommands) HasRemoteModel(context.Context, string) (bool, error) { + f.rec("HasRemoteModel") + return true, nil +} +func (f *fakeCommands) InstallTimeout() time.Duration { f.rec("InstallTimeout"); return time.Minute } +func (f *fakeCommands) DeleteModelFiles(string) error { f.rec("DeleteModelFiles"); return nil } +func (f *fakeCommands) InstallBackendForce(string, string, string, string, string, string, int, string, func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + f.rec("InstallBackendForce") + return &workerctl.BackendInstallReply{}, nil +} + +// fakeFiles is a carrier.FileCarrier that only records the call. +type fakeFiles struct{ recorder } + +var _ carrier.FileCarrier = (*fakeFiles)(nil) + +func (f *fakeFiles) EnsureRemote(context.Context, string, string, string) (string, error) { + f.rec("EnsureRemote") + return "", nil +} + +func (f *fakeFiles) FetchRemote(context.Context, string, string, string) error { + f.rec("FetchRemote") + return nil +} + +func (f *fakeFiles) FetchRemoteByKey(context.Context, string, string, string) error { + f.rec("FetchRemoteByKey") + return nil +} + +func (f *fakeFiles) AllocRemoteTemp(context.Context, string) (string, error) { + f.rec("AllocRemoteTemp") + return "", nil +} + +func (f *fakeFiles) StageRemoteToStore(context.Context, string, string, string) error { + f.rec("StageRemoteToStore") + return nil +} + +func (f *fakeFiles) ReleaseRemote(context.Context, string, string) error { + f.rec("ReleaseRemote") + return nil +} + +func (f *fakeFiles) ListRemoteDir(context.Context, string, string) ([]string, error) { + f.rec("ListRemoteDir") + return nil, nil +} + +func (f *fakeFiles) ReleaseRemoteRequest(context.Context, string, string, []string) error { + f.rec("ReleaseRemoteRequest") + return nil +} + +type fakeQueue struct{ recorder } + +func (f *fakeQueue) Enqueue(context.Context, messaging.WorkKind, any) error { + f.rec("Enqueue") + return nil +} + +type fakeClients struct { + recorder + gotNode, gotAddr string + gotParallel bool +} + +func (f *fakeClients) NewClient(nodeID, address string, parallel bool) grpc.Backend { + f.rec("NewClient") + f.gotNode, f.gotAddr, f.gotParallel = nodeID, address, parallel + return nil +} + +type fakeAgents struct{ recorder } + +func (f *fakeAgents) ExecuteMCPTool(context.Context, mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error) { + f.rec("ExecuteMCPTool") + return &mcpRemote.MCPToolResponse{}, nil +} + +func (f *fakeAgents) DiscoverMCPTools(context.Context, mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error) { + f.rec("DiscoverMCPTools") + return &mcpRemote.MCPDiscoveryResponse{}, nil +} + +// fakeCarrier is one complete carrier set with every member recorded. +type fakeCarrier struct { + set *carrier.Set + bus *testutil.FakeBus + queue *fakeQueue + commands *fakeCommands + files *fakeFiles + clients *fakeClients + agents *fakeAgents + dialed *recorder +} + +// newFakeCarrier builds a set named name. The two sets of a spec share a +// carrier name only when the spec says so. +func newFakeCarrier(name cluster.Carrier, epoch int64) *fakeCarrier { + f := &fakeCarrier{ + bus: testutil.NewFakeBus(), + queue: &fakeQueue{}, + commands: &fakeCommands{}, + files: &fakeFiles{}, + clients: &fakeClients{}, + agents: &fakeAgents{}, + dialed: &recorder{}, + } + f.set = &carrier.Set{ + Name: name, + Epoch: epoch, + Broadcaster: f.bus, + OnReconnect: f.bus.OnReconnect, + WorkQueue: f.queue, + Commands: f.commands, + Files: f.files, + Clients: f.clients, + Dialer: func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + return func(context.Context, string, string) (net.Conn, error) { + f.dialed.rec("dial:" + nodeID) + return nil, nil + } + }, + Agents: f.agents, + } + return f +} diff --git a/core/services/carrier/holders.go b/core/services/carrier/holders.go new file mode 100644 index 000000000000..2388d9c7c7da --- /dev/null +++ b/core/services/carrier/holders.go @@ -0,0 +1,192 @@ +package carrier + +import ( + "context" + "net" + "sync/atomic" + "time" + + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +// Every holder below does one thing per call: load the pointer once and call +// the same method on that set. Loading once pins a multi-step operation of the +// carrier to one set, and a call that is already running is never moved. The +// holders take no lock and allocate nothing of their own. + +// Commands forwards the control verbs to the current set. +type Commands struct{ cur *atomic.Pointer[Set] } + +var _ nodes.NodeControl = (*Commands)(nil) + +// NewCommands returns a holder over cur, which must name a set. +func NewCommands(cur *atomic.Pointer[Set]) *Commands { return &Commands{cur: cur} } + +func (h *Commands) InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + return h.cur.Load().Commands.InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) +} + +func (h *Commands) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error) { + return h.cur.Load().Commands.UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) +} + +func (h *Commands) DeleteBackend(nodeID, backendName string) (*workerctl.BackendDeleteReply, error) { + return h.cur.Load().Commands.DeleteBackend(nodeID, backendName) +} + +func (h *Commands) ListBackends(nodeID string) (*workerctl.BackendListReply, error) { + return h.cur.Load().Commands.ListBackends(nodeID) +} + +func (h *Commands) StopBackend(nodeID, backend string) error { + return h.cur.Load().Commands.StopBackend(nodeID, backend) +} + +func (h *Commands) UnloadModelOnNode(nodeID, modelName string) error { + return h.cur.Load().Commands.UnloadModelOnNode(nodeID, modelName) +} + +func (h *Commands) PingNode(nodeID string) error { return h.cur.Load().Commands.PingNode(nodeID) } + +func (h *Commands) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + return h.cur.Load().Commands.InstallBackendOp(nodeID, backendType, modelID, galleriesJSON, replicaIndex, opID, operationID, deadline, onProgress) +} + +func (h *Commands) StopLoadOperation(ctx context.Context, nodeID string, req workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { + return h.cur.Load().Commands.StopLoadOperation(ctx, nodeID, req) +} + +func (h *Commands) OperationControl(nodeID string, req workerctl.OperationRequest) (*workerctl.OperationReply, error) { + return h.cur.Load().Commands.OperationControl(nodeID, req) +} + +func (h *Commands) UnloadReplica(nodeID string, replica nodes.NodeModel) error { + return h.cur.Load().Commands.UnloadReplica(nodeID, replica) +} + +func (h *Commands) StopModelReplica(ctx context.Context, nodeID string, replica nodes.NodeModel, force bool) (workerctl.ModelStopReply, error) { + return h.cur.Load().Commands.StopModelReplica(ctx, nodeID, replica, force) +} + +func (h *Commands) ListRunningModels(nodeID string) (*workerctl.ModelsRunningReply, error) { + return h.cur.Load().Commands.ListRunningModels(nodeID) +} + +func (h *Commands) UnloadRemoteModel(modelName string) error { + return h.cur.Load().Commands.UnloadRemoteModel(modelName) +} + +func (h *Commands) UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error { + return h.cur.Load().Commands.UnloadRemoteModelContext(ctx, modelName, force) +} + +func (h *Commands) HasRemoteModel(ctx context.Context, modelName string) (bool, error) { + return h.cur.Load().Commands.HasRemoteModel(ctx, modelName) +} + +func (h *Commands) InstallTimeout() time.Duration { return h.cur.Load().Commands.InstallTimeout() } + +func (h *Commands) DeleteModelFiles(modelName string) error { + return h.cur.Load().Commands.DeleteModelFiles(modelName) +} + +func (h *Commands) InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + return h.cur.Load().Commands.InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) +} + +// Files forwards file staging to the current set. +type Files struct{ cur *atomic.Pointer[Set] } + +var _ FileCarrier = (*Files)(nil) + +// NewFiles returns a holder over cur, which must name a set. +func NewFiles(cur *atomic.Pointer[Set]) *Files { return &Files{cur: cur} } + +func (h *Files) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) { + return h.cur.Load().Files.EnsureRemote(ctx, nodeID, localPath, key) +} + +func (h *Files) FetchRemote(ctx context.Context, nodeID, remotePath, localDst string) error { + return h.cur.Load().Files.FetchRemote(ctx, nodeID, remotePath, localDst) +} + +func (h *Files) FetchRemoteByKey(ctx context.Context, nodeID, key, localDst string) error { + return h.cur.Load().Files.FetchRemoteByKey(ctx, nodeID, key, localDst) +} + +func (h *Files) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) { + return h.cur.Load().Files.AllocRemoteTemp(ctx, nodeID) +} + +func (h *Files) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error { + return h.cur.Load().Files.StageRemoteToStore(ctx, nodeID, remotePath, key) +} + +func (h *Files) ReleaseRemote(ctx context.Context, nodeID, key string) error { + return h.cur.Load().Files.ReleaseRemote(ctx, nodeID, key) +} + +func (h *Files) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) { + return h.cur.Load().Files.ListRemoteDir(ctx, nodeID, keyPrefix) +} + +func (h *Files) ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error { + return h.cur.Load().Files.ReleaseRemoteRequest(ctx, nodeID, requestID, keys) +} + +// WorkQueue forwards Enqueue to the current set. Producers always use the +// active carrier. +type WorkQueue struct{ cur *atomic.Pointer[Set] } + +var _ messaging.WorkQueue = (*WorkQueue)(nil) + +// NewWorkQueue returns a holder over cur, which must name a set. +func NewWorkQueue(cur *atomic.Pointer[Set]) *WorkQueue { return &WorkQueue{cur: cur} } + +func (h *WorkQueue) Enqueue(ctx context.Context, kind messaging.WorkKind, payload any) error { + return h.cur.Load().WorkQueue.Enqueue(ctx, kind, payload) +} + +// Clients forwards backend client construction to the current set. A client +// built before a swap keeps the connection it has; the code that caches +// clients drops them when the carrier changes. +type Clients struct{ cur *atomic.Pointer[Set] } + +var _ nodes.BackendClientFactory = (*Clients)(nil) + +// NewClients returns a holder over cur, which must name a set. +func NewClients(cur *atomic.Pointer[Set]) *Clients { return &Clients{cur: cur} } + +func (h *Clients) NewClient(nodeID, address string, parallel bool) grpc.Backend { + return h.cur.Load().Clients.NewClient(nodeID, address, parallel) +} + +// NewWorkerDialer returns a nodes.WorkerNetDialerFor that picks the set at the +// moment of each dial, not when the dial function is created. An HTTP client +// keeps its dial function for as long as it lives, so resolving earlier would +// pin it to the carrier it was created on. +func NewWorkerDialer(cur *atomic.Pointer[Set]) nodes.WorkerNetDialerFor { + return func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, network, addr string) (net.Conn, error) { + return cur.Load().Dialer(nodeID)(ctx, network, addr) + } + } +} + +// Agents forwards MCP requests to the agent workers of the current set. +type Agents struct{ cur *atomic.Pointer[Set] } + +// NewAgents returns a holder over cur, which must name a set. +func NewAgents(cur *atomic.Pointer[Set]) *Agents { return &Agents{cur: cur} } + +func (h *Agents) ExecuteMCPTool(ctx context.Context, req mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error) { + return h.cur.Load().Agents.ExecuteMCPTool(ctx, req) +} + +func (h *Agents) DiscoverMCPTools(ctx context.Context, req mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error) { + return h.cur.Load().Agents.DiscoverMCPTools(ctx, req) +} diff --git a/core/services/carrier/holders_test.go b/core/services/carrier/holders_test.go new file mode 100644 index 000000000000..d66a7df93961 --- /dev/null +++ b/core/services/carrier/holders_test.go @@ -0,0 +1,310 @@ +package carrier_test + +import ( + "context" + "net" + "sync/atomic" + "testing" + "time" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" + grpc "github.com/mudler/LocalAI/pkg/grpc" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// callCase calls one method of a holder. +type callCase[H any] struct { + method string + call func(H) +} + +var _ = Describe("Set", func() { + It("accepts a complete set", func() { + Expect(newFakeCarrier(cluster.CarrierNATS, 1).set.Validate()).To(Succeed()) + }) + + It("names the first member that is missing", func() { + for name, mutate := range map[string]func(*carrier.Set){ + "Broadcaster": func(s *carrier.Set) { s.Broadcaster = nil }, + "WorkQueue": func(s *carrier.Set) { s.WorkQueue = nil }, + "Commands": func(s *carrier.Set) { s.Commands = nil }, + "Files": func(s *carrier.Set) { s.Files = nil }, + "Clients": func(s *carrier.Set) { s.Clients = nil }, + "Dialer": func(s *carrier.Set) { s.Dialer = nil }, + "Agents": func(s *carrier.Set) { s.Agents = nil }, + } { + s := newFakeCarrier(cluster.CarrierNATS, 1).set + mutate(s) + Expect(s.Validate()).To(MatchError(ContainSubstring(name)), name) + } + }) + + It("refuses a carrier name the cluster does not know", func() { + s := newFakeCarrier("carrier-pigeon", 1).set + Expect(s.Validate()).To(MatchError(cluster.ErrInvalidCarrier)) + }) +}) + +var _ = Describe("Forwarding holders", func() { + var ( + a, b *fakeCarrier + cur atomic.Pointer[carrier.Set] + ) + + BeforeEach(func() { + a = newFakeCarrier(cluster.CarrierNATS, 1) + b = newFakeCarrier(cluster.CarrierTunnel, 2) + cur = atomic.Pointer[carrier.Set]{} + cur.Store(a.set) + }) + + ctx := context.Background() + + Describe("Commands", func() { + var progress = func(workerctl.BackendInstallProgressEvent) {} + cases := []callCase[*carrier.Commands]{ + {"InstallBackend", func(h *carrier.Commands) { + _, _ = h.InstallBackend("n", "t", "m", "g", "u", "name", "alias", 1, "op", progress) + }}, + {"UpgradeBackend", func(h *carrier.Commands) { + _, _ = h.UpgradeBackend("n", "t", "g", "u", "name", "alias", 1, "op", progress) + }}, + {"DeleteBackend", func(h *carrier.Commands) { _, _ = h.DeleteBackend("n", "b") }}, + {"ListBackends", func(h *carrier.Commands) { _, _ = h.ListBackends("n") }}, + {"StopBackend", func(h *carrier.Commands) { _ = h.StopBackend("n", "b") }}, + {"UnloadModelOnNode", func(h *carrier.Commands) { _ = h.UnloadModelOnNode("n", "m") }}, + {"PingNode", func(h *carrier.Commands) { _ = h.PingNode("n") }}, + {"InstallBackendOp", func(h *carrier.Commands) { + _, _ = h.InstallBackendOp("n", "t", "m", "g", 1, "op", "oid", time.Second, progress) + }}, + {"StopLoadOperation", func(h *carrier.Commands) { + _, _ = h.StopLoadOperation(ctx, "n", workerctl.ModelStopRequest{}) + }}, + {"OperationControl", func(h *carrier.Commands) { _, _ = h.OperationControl("n", workerctl.OperationRequest{}) }}, + {"UnloadReplica", func(h *carrier.Commands) { _ = h.UnloadReplica("n", nodes.NodeModel{}) }}, + {"StopModelReplica", func(h *carrier.Commands) { + _, _ = h.StopModelReplica(ctx, "n", nodes.NodeModel{}, true) + }}, + {"ListRunningModels", func(h *carrier.Commands) { _, _ = h.ListRunningModels("n") }}, + {"UnloadRemoteModel", func(h *carrier.Commands) { _ = h.UnloadRemoteModel("m") }}, + {"UnloadRemoteModelContext", func(h *carrier.Commands) { _ = h.UnloadRemoteModelContext(ctx, "m", true) }}, + {"HasRemoteModel", func(h *carrier.Commands) { _, _ = h.HasRemoteModel(ctx, "m") }}, + {"InstallTimeout", func(h *carrier.Commands) { _ = h.InstallTimeout() }}, + {"DeleteModelFiles", func(h *carrier.Commands) { _ = h.DeleteModelFiles("m") }}, + {"InstallBackendForce", func(h *carrier.Commands) { + _, _ = h.InstallBackendForce("n", "t", "g", "u", "name", "alias", 1, "op", progress) + }}, + } + + It("covers every method of nodes.NodeControl", func() { + // The compile-time assertion proves the holder has each method; this + // table proves each one is forwarded to the method of the same name. + Expect(cases).To(HaveLen(19)) + }) + + for _, tc := range cases { + It("forwards "+tc.method+" to the carrier the pointer names", func() { + h := carrier.NewCommands(&cur) + tc.call(h) + Expect(a.commands.seen()).To(Equal([]string{tc.method})) + Expect(b.commands.seen()).To(BeEmpty()) + + cur.Store(b.set) + tc.call(h) + Expect(a.commands.seen()).To(Equal([]string{tc.method})) + Expect(b.commands.seen()).To(Equal([]string{tc.method})) + }) + } + + It("passes the arguments and the reply through untouched", func() { + var got []string + fake := &argCommands{fakeCommands: a.commands, got: &got} + a.set.Commands = fake + h := carrier.NewCommands(&cur) + + reply, err := h.OperationControl("node-7", workerctl.OperationRequest{Renew: []string{"op-9"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Renewed).To(Equal([]string{"op-9"})) + Expect(got).To(Equal([]string{"node-7/op-9"})) + }) + }) + + Describe("Files", func() { + cases := []callCase[*carrier.Files]{ + {"EnsureRemote", func(h *carrier.Files) { _, _ = h.EnsureRemote(ctx, "n", "l", "k") }}, + {"FetchRemote", func(h *carrier.Files) { _ = h.FetchRemote(ctx, "n", "r", "d") }}, + {"FetchRemoteByKey", func(h *carrier.Files) { _ = h.FetchRemoteByKey(ctx, "n", "k", "d") }}, + {"AllocRemoteTemp", func(h *carrier.Files) { _, _ = h.AllocRemoteTemp(ctx, "n") }}, + {"StageRemoteToStore", func(h *carrier.Files) { _ = h.StageRemoteToStore(ctx, "n", "r", "k") }}, + {"ReleaseRemote", func(h *carrier.Files) { _ = h.ReleaseRemote(ctx, "n", "k") }}, + {"ListRemoteDir", func(h *carrier.Files) { _, _ = h.ListRemoteDir(ctx, "n", "p") }}, + {"ReleaseRemoteRequest", func(h *carrier.Files) { _ = h.ReleaseRemoteRequest(ctx, "n", "req", []string{"k"}) }}, + } + + for _, tc := range cases { + It("forwards "+tc.method+" to the carrier the pointer names", func() { + h := carrier.NewFiles(&cur) + tc.call(h) + Expect(a.files.seen()).To(Equal([]string{tc.method})) + cur.Store(b.set) + tc.call(h) + Expect(b.files.seen()).To(Equal([]string{tc.method})) + Expect(a.files.seen()).To(HaveLen(1)) + }) + } + + It("is the batch release the file staging client looks for", func() { + var stager nodes.FileStager = carrier.NewFiles(&cur) + _, ok := stager.(nodes.RequestFileReleaser) + Expect(ok).To(BeTrue(), "a holder that hid the extension would silently turn batch release into one call per key") + }) + }) + + Describe("WorkQueue", func() { + It("enqueues on the carrier the pointer names", func() { + h := carrier.NewWorkQueue(&cur) + Expect(h.Enqueue(ctx, messaging.WorkTask, "x")).To(Succeed()) + cur.Store(b.set) + Expect(h.Enqueue(ctx, messaging.WorkMCPCI, "y")).To(Succeed()) + Expect(a.queue.seen()).To(HaveLen(1)) + Expect(b.queue.seen()).To(HaveLen(1)) + }) + }) + + Describe("Clients", func() { + It("builds the client with the factory the pointer names", func() { + h := carrier.NewClients(&cur) + _ = h.NewClient("node-1", "10.0.0.1:50051", true) + Expect(a.clients.gotNode).To(Equal("node-1")) + Expect(a.clients.gotAddr).To(Equal("10.0.0.1:50051")) + Expect(a.clients.gotParallel).To(BeTrue()) + + cur.Store(b.set) + _ = h.NewClient("node-2", "x", false) + Expect(a.clients.seen()).To(HaveLen(1)) + Expect(b.clients.gotNode).To(Equal("node-2")) + }) + }) + + Describe("WorkerDialer", func() { + It("dials with the carrier the pointer names at the moment of each dial", func() { + dialFor := carrier.NewWorkerDialer(&cur) + dial := dialFor("node-1") // a client keeps this function for good + + _, _ = dial(ctx, "tcp", "x:1") + cur.Store(b.set) + _, _ = dial(ctx, "tcp", "x:1") + + Expect(a.dialed.seen()).To(Equal([]string{"dial:node-1"})) + Expect(b.dialed.seen()).To(Equal([]string{"dial:node-1"})) + }) + }) + + Describe("Agents", func() { + It("forwards both verbs to the carrier the pointer names", func() { + h := carrier.NewAgents(&cur) + _, _ = h.ExecuteMCPTool(ctx, mcpRemote.MCPToolRequest{}) + _, _ = h.DiscoverMCPTools(ctx, mcpRemote.MCPDiscoveryRequest{}) + cur.Store(b.set) + _, _ = h.ExecuteMCPTool(ctx, mcpRemote.MCPToolRequest{}) + Expect(a.agents.seen()).To(Equal([]string{"ExecuteMCPTool", "DiscoverMCPTools"})) + Expect(b.agents.seen()).To(Equal([]string{"ExecuteMCPTool"})) + }) + }) + + Describe("a call in flight", func() { + It("finishes on the carrier it started on, and a new call goes to the new one", func() { + gate := make(chan struct{}) + entered := make(chan struct{}, 1) + a.set.Commands = &gatedCommands{fakeCommands: a.commands, gate: gate, entered: entered} + h := carrier.NewCommands(&cur) + + done := make(chan error, 1) + go func() { + _, err := h.InstallBackendOp("n", "t", "m", "g", 0, "op", "oid", time.Minute, nil) + done <- err + }() + Eventually(entered).Should(Receive()) + + cur.Store(b.set) + Expect(h.PingNode("n")).To(Succeed()) + Expect(b.commands.seen()).To(Equal([]string{"PingNode"})) + + close(gate) + Eventually(done, time.Second).Should(Receive(BeNil())) + Expect(a.commands.seen()).To(Equal([]string{"InstallBackendOp"})) + }) + }) + + Describe("hot path", func() { + It("adds no allocation to a forwarded call", func() { + h := carrier.NewCommands(&cur) + wq := carrier.NewWorkQueue(&cur) + cl := carrier.NewClients(&cur) + dial := carrier.NewWorkerDialer(&cur)("node-1") + + cur.Store(newQuietCarrier().set) + var payload any = "payload" + + Expect(testing.AllocsPerRun(1000, func() { _ = h.PingNode("node-1") })).To(BeZero()) + Expect(testing.AllocsPerRun(1000, func() { _ = wq.Enqueue(ctx, messaging.WorkTask, payload) })).To(BeZero()) + Expect(testing.AllocsPerRun(1000, func() { _ = cl.NewClient("node-1", "x:1", false) })).To(BeZero()) + Expect(testing.AllocsPerRun(1000, func() { _, _ = dial(ctx, "tcp", "x:1") })).To(BeZero()) + }) + }) +}) + +// argCommands answers OperationControl with a reply that records its arguments. +type argCommands struct { + *fakeCommands + got *[]string +} + +func (c *argCommands) OperationControl(nodeID string, req workerctl.OperationRequest) (*workerctl.OperationReply, error) { + *c.got = append(*c.got, nodeID+"/"+req.Renew[0]) + return &workerctl.OperationReply{Renewed: req.Renew}, nil +} + +// gatedCommands blocks InstallBackendOp until gate closes. +type gatedCommands struct { + *fakeCommands + gate chan struct{} + entered chan struct{} +} + +func (c *gatedCommands) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + c.entered <- struct{}{} + <-c.gate + return c.fakeCommands.InstallBackendOp(nodeID, backendType, modelID, galleriesJSON, replicaIndex, opID, operationID, deadline, onProgress) +} + +// quietCommands, quietQueue and quietClients answer without recording, so a +// spec can count the allocations the holder itself adds. +type quietCommands struct{ *fakeCommands } + +func (quietCommands) PingNode(string) error { return nil } + +type quietQueue struct{} + +func (quietQueue) Enqueue(context.Context, messaging.WorkKind, any) error { return nil } + +type quietClients struct{} + +func (quietClients) NewClient(string, string, bool) grpc.Backend { return nil } + +func newQuietCarrier() *fakeCarrier { + f := newFakeCarrier(cluster.CarrierNATS, 9) + f.set.Commands = quietCommands{f.commands} + f.set.WorkQueue = quietQueue{} + f.set.Clients = quietClients{} + f.set.Dialer = func(string) func(context.Context, string, string) (net.Conn, error) { return quietDial } + return f +} + +func quietDial(context.Context, string, string) (net.Conn, error) { return nil, nil } diff --git a/core/services/carrier/import_boundary_test.go b/core/services/carrier/import_boundary_test.go new file mode 100644 index 000000000000..98f30aad6516 --- /dev/null +++ b/core/services/carrier/import_boundary_test.go @@ -0,0 +1,184 @@ +package carrier_test + +import ( + "go/parser" + "go/token" + "io/fs" + "os" + "path/filepath" + "slices" + "strconv" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Only the files of a carrier may import the library of that carrier. The code +// above the seams, and the holders in this package, name interfaces only, so a +// second carrier is a new implementation and not a change to its callers. +// +// The scan reads the import lines of every non-test Go file under the module +// root. It does not use `go list`: that skips a file whose build constraints do +// not match the host, and an import behind a constraint would pass unseen. It +// does not use a linter rule either: the repository does not run depguard, and +// a spec fails in the same place every other spec does. +// +// Test files are exempt. They start servers and mint credentials with the +// library on purpose. + +const natsLibraries = "github.com/nats-io/" + +// natsAllowlist is every non-test file or package directory allowed to import +// natsLibraries, as a slash path from the module root. Add to it only for a +// file that is part of the NATS carrier. A new carrier library gets its own +// list in the slice that adds the carrier. +var natsAllowlist = []string{ + "core/services/messaging/client.go", + "core/services/messaging/tls.go", + "core/services/nodes/control_nats.go", + "pkg/natsauth/", // credential minting and decoding for the NATS carrier +} + +// importers returns the slash paths, from root, of the non-test Go files that +// import a path with the given prefix. It skips hidden directories, +// node_modules and nested modules. +func importers(root, prefix string) ([]string, error) { + var found []string + fset := token.NewFileSet() + err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + name := d.Name() + if d.IsDir() { + if path != root { + if strings.HasPrefix(name, ".") || name == "node_modules" { + return filepath.SkipDir + } + if _, err := os.Stat(filepath.Join(path, "go.mod")); err == nil { + return filepath.SkipDir + } + } + return nil + } + if !strings.HasSuffix(name, ".go") || strings.HasSuffix(name, "_test.go") { + return nil + } + f, err := parser.ParseFile(fset, path, nil, parser.ImportsOnly) + if err != nil { + return err + } + for _, imp := range f.Imports { + p, err := strconv.Unquote(imp.Path.Value) + if err != nil { + return err + } + if strings.HasPrefix(p, prefix) { + rel, err := filepath.Rel(root, path) + if err != nil { + return err + } + found = append(found, filepath.ToSlash(rel)) + return nil + } + } + return nil + }) + slices.Sort(found) + return found, err +} + +func allowed(file string, allowlist []string) bool { + for _, a := range allowlist { + if file == a || (strings.HasSuffix(a, "/") && strings.HasPrefix(file, a)) { + return true + } + } + return false +} + +// moduleRoot walks up from the working directory of the spec to go.mod. +func moduleRoot() string { + dir, err := os.Getwd() + Expect(err).ToNot(HaveOccurred()) + for { + if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil { + return dir + } + parent := filepath.Dir(dir) + Expect(parent).ToNot(Equal(dir), "no go.mod above the spec") + dir = parent + } +} + +var _ = Describe("Import boundary of the NATS libraries", func() { + It("lets only the files of the NATS carrier import them", func() { + files, err := importers(moduleRoot(), natsLibraries) + Expect(err).ToNot(HaveOccurred()) + + var stray []string + for _, f := range files { + if !allowed(f, natsAllowlist) { + stray = append(stray, f) + } + } + Expect(stray).To(BeEmpty(), "these files import %s but are not in natsAllowlist; code above the seams must use messaging and nodes interfaces", natsLibraries) + }) + + It("keeps the allowlist honest: every entry still imports the libraries", func() { + files, err := importers(moduleRoot(), natsLibraries) + Expect(err).ToNot(HaveOccurred()) + + for _, a := range natsAllowlist { + used := slices.ContainsFunc(files, func(f string) bool { return allowed(f, []string{a}) }) + Expect(used).To(BeTrue(), "%s no longer imports %s; remove it from natsAllowlist", a, natsLibraries) + } + }) + + Describe("the scan", func() { + var root string + + write := func(rel, body string) { + path := filepath.Join(root, filepath.FromSlash(rel)) + Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed()) + Expect(os.WriteFile(path, []byte(body), 0o600)).To(Succeed()) + } + + BeforeEach(func() { root = GinkgoT().TempDir() }) + + It("finds an import in any non-test file, including one behind a build constraint", func() { + write("a/plain.go", "package a\nimport _ \"github.com/nats-io/nats.go\"\n") + write("a/tagged.go", "//go:build never\n\npackage a\nimport _ \"github.com/nats-io/nkeys\"\n") + write("b/clean.go", "package b\nimport _ \"fmt\"\n") + + files, err := importers(root, natsLibraries) + Expect(err).ToNot(HaveOccurred()) + Expect(files).To(Equal([]string{"a/plain.go", "a/tagged.go"})) + }) + + It("ignores test files, hidden directories and nested modules", func() { + write("a/x_test.go", "package a\nimport _ \"github.com/nats-io/nats.go\"\n") + write(".hidden/h.go", "package h\nimport _ \"github.com/nats-io/nats.go\"\n") + write("nested/go.mod", "module nested\n") + write("nested/n.go", "package n\nimport _ \"github.com/nats-io/nats.go\"\n") + + files, err := importers(root, natsLibraries) + Expect(err).ToNot(HaveOccurred()) + Expect(files).To(BeEmpty()) + }) + + It("reports a file it cannot parse", func() { + write("a/broken.go", "package a\nimport (\n") + _, err := importers(root, natsLibraries) + Expect(err).To(HaveOccurred()) + }) + + It("matches an allowlist entry as a file or as a directory", func() { + Expect(allowed("pkg/natsauth/mint.go", []string{"pkg/natsauth/"})).To(BeTrue()) + Expect(allowed("pkg/natsauthx/mint.go", []string{"pkg/natsauth/"})).To(BeFalse()) + Expect(allowed("core/a.go", []string{"core/a.go"})).To(BeTrue()) + Expect(allowed("core/ab.go", []string{"core/a.go"})).To(BeFalse()) + }) + }) +}) diff --git a/core/services/carrier/nats.go b/core/services/carrier/nats.go new file mode 100644 index 000000000000..04415c0bcee8 --- /dev/null +++ b/core/services/carrier/nats.go @@ -0,0 +1,78 @@ +package carrier + +import ( + "fmt" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/storage" +) + +// NATSClient is the connection a NATS set is built over. *messaging.Client +// satisfies it. +type NATSClient interface { + messaging.MessagingClient + OnReconnect(func()) +} + +// NATSOptions is what NewNATSSet needs besides the connection. +type NATSOptions struct { + Client NATSClient + // Epoch is the epoch of the cluster row the set is built for. + Epoch int64 + + // Registry finds the nodes that hold a model, for the verbs that fan out. + Registry nodes.ModelLocator + // InstallTimeout and UpgradeTimeout bound the install and upgrade requests. + InstallTimeout, UpgradeTimeout time.Duration + + // Token is the registration token. It is the bearer credential of the + // gRPC clients and of the file transfer to a worker. + Token string + + // S3Staging selects the file stager: shared object storage with a NATS + // request to the worker when true, a direct HTTP transfer to the worker + // when false. + S3Staging bool + // FileManager is the object storage client of the S3 stager. + FileManager *storage.FileManager + // HTTPAddrFor returns the address of a worker's file transfer server, for + // the HTTP stager. + HTTPAddrFor func(nodeID string) (string, error) +} + +// NewNATSSet builds the set of seam implementations that run over NATS for +// the control plane and over a direct dial to the worker for gRPC and HTTP. +func NewNATSSet(o NATSOptions) (*Set, error) { + if o.Client == nil { + return nil, fmt.Errorf("NATS set needs a connection") + } + dial := nodes.DirectWorkerNetDialer() + + var files FileCarrier + if o.S3Staging { + files = nodes.NewS3NATSFileStager(o.FileManager, o.Client) + } else { + files = nodes.NewHTTPFileStager(o.HTTPAddrFor, o.Token, dial) + } + + set := &Set{ + Name: cluster.CarrierNATS, + Epoch: o.Epoch, + Broadcaster: o.Client, + OnReconnect: o.Client.OnReconnect, + WorkQueue: messaging.NewNATSWorkQueue(o.Client), + Commands: nodes.NewRemoteUnloaderAdapter(o.Registry, o.Client, o.InstallTimeout, o.UpgradeTimeout), + Files: files, + Clients: nodes.NewTokenClientFactory(o.Token), + Dialer: dial, + Agents: nodes.NewNATSAgentControl(o.Client), + Close: o.Client.Close, + } + if err := set.Validate(); err != nil { + return nil, err + } + return set, nil +} diff --git a/core/services/carrier/nats_test.go b/core/services/carrier/nats_test.go new file mode 100644 index 000000000000..b8b285677383 --- /dev/null +++ b/core/services/carrier/nats_test.go @@ -0,0 +1,95 @@ +package carrier_test + +import ( + "sync/atomic" + "time" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// closingBus counts Close calls on a FakeBus. +type closingBus struct { + *testutil.FakeBus + closed atomic.Int32 +} + +func (c *closingBus) Close() { c.closed.Add(1) } + +// noModels is a ModelLocator that knows no node. +type noModels struct{ nodes.ModelLocator } + +var _ = Describe("NewNATSSet", func() { + var bus *closingBus + + BeforeEach(func() { + bus = &closingBus{FakeBus: testutil.NewFakeBus()} + }) + + opts := func() carrier.NATSOptions { + return carrier.NATSOptions{ + Client: bus, + Epoch: 7, + Registry: noModels{}, + InstallTimeout: 3 * time.Minute, + UpgradeTimeout: 4 * time.Minute, + Token: "token", + HTTPAddrFor: func(string) (string, error) { return "127.0.0.1:1", nil }, + } + } + + It("builds a complete set for the NATS carrier", func() { + set, err := carrier.NewNATSSet(opts()) + Expect(err).ToNot(HaveOccurred()) + Expect(set.Validate()).To(Succeed()) + Expect(set.Name).To(Equal(cluster.CarrierNATS)) + Expect(set.Epoch).To(Equal(int64(7))) + Expect(set.Commands.InstallTimeout()).To(Equal(3 * time.Minute)) + }) + + It("publishes and enqueues over the connection it was given", func() { + set, err := carrier.NewNATSSet(opts()) + Expect(err).ToNot(HaveOccurred()) + Expect(set.Broadcaster.Publish("nodes.x", "hello")).To(Succeed()) + Expect(bus.PublishCount("nodes.x")).To(Equal(1)) + }) + + It("stages files over HTTP unless shared storage is configured", func() { + set, err := carrier.NewNATSSet(opts()) + Expect(err).ToNot(HaveOccurred()) + Expect(set.Files).To(BeAssignableToTypeOf(&nodes.HTTPFileStager{})) + + o := opts() + o.S3Staging = true + set, err = carrier.NewNATSSet(o) + Expect(err).ToNot(HaveOccurred()) + Expect(set.Files).To(BeAssignableToTypeOf(&nodes.S3NATSFileStager{})) + }) + + It("forwards the reconnect hook of the connection", func() { + set, err := carrier.NewNATSSet(opts()) + Expect(err).ToNot(HaveOccurred()) + var n atomic.Int32 + set.OnReconnect(func() { n.Add(1) }) + bus.TriggerReconnect() + Expect(n.Load()).To(Equal(int32(1))) + }) + + It("closes the connection when the set is closed", func() { + set, err := carrier.NewNATSSet(opts()) + Expect(err).ToNot(HaveOccurred()) + set.Close() + Expect(bus.closed.Load()).To(Equal(int32(1))) + }) + + It("refuses to build without a connection", func() { + o := opts() + o.Client = nil + _, err := carrier.NewNATSSet(o) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/carrier/set.go b/core/services/carrier/set.go new file mode 100644 index 000000000000..2a5cade12124 --- /dev/null +++ b/core/services/carrier/set.go @@ -0,0 +1,81 @@ +// Package carrier makes the transport of a distributed deployment replaceable +// at runtime. +// +// Every seam that crosses a process boundary (fan-out, queues, control verbs, +// file staging, backend dials, agent RPC) is an interface. A Set is one +// carrier's implementation of all of them. The holders in this package +// implement the same interfaces and forward each call to the Set that an +// atomic pointer names at that moment, so code that was built against an +// interface keeps working while the pointer moves. +// +// This package names interfaces only. The code that builds a Set for a +// carrier is the one place that names the carrier (see nats.go). +package carrier + +import ( + "errors" + "fmt" + + mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" +) + +// FileCarrier is a file stager with the batch release the file staging client +// looks for. It is one interface so that a carrier cannot leave the extension +// out, and the holder cannot hide it. +type FileCarrier interface { + nodes.FileStager + nodes.RequestFileReleaser +} + +// Set is the implementation of every seam for one carrier. Build it completely +// and call Validate before storing it where holders can see it. +type Set struct { + // Name is the carrier this set belongs to. + Name cluster.Carrier + // Epoch is the epoch of the cluster row the set was built for. + Epoch int64 + + Broadcaster messaging.Broadcaster + // OnReconnect registers a callback the carrier runs after it recovers a + // lost connection. It is nil for a carrier with no such event. + OnReconnect func(func()) + WorkQueue messaging.WorkQueue + Commands nodes.NodeControl + Files FileCarrier + Clients nodes.BackendClientFactory + Dialer nodes.WorkerNetDialerFor + Agents mcpTools.AgentControl + + // Close releases the carrier's connections. It may be nil. + Close func() +} + +// Validate reports the first member that is missing. +func (s *Set) Validate() error { + if s == nil { + return errors.New("carrier set is nil") + } + if s.Name != cluster.CarrierNATS && s.Name != cluster.CarrierTunnel { + return fmt.Errorf("%w: %q", cluster.ErrInvalidCarrier, s.Name) + } + for _, m := range []struct { + name string + ok bool + }{ + {"Broadcaster", s.Broadcaster != nil}, + {"WorkQueue", s.WorkQueue != nil}, + {"Commands", s.Commands != nil}, + {"Files", s.Files != nil}, + {"Clients", s.Clients != nil}, + {"Dialer", s.Dialer != nil}, + {"Agents", s.Agents != nil}, + } { + if !m.ok { + return fmt.Errorf("carrier set %q has no %s", s.Name, m.name) + } + } + return nil +} From 7ad74d305cc69d4ebdd0e92049bfca50120e8a59 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:18:54 +0000 Subject: [PATCH 004/162] refactor(distributed): build the distributed services over carrier holders Read the cluster carrier row at startup, or insert it when it is absent, and build the set it names. An existing deployment has a NATS URL, so it seeds nats and runs as before. A row that names a carrier this build does not run stops the start with a message, and does not fall back. DistributedServices no longer holds the NATS client. The router, the health monitor, the reconciler, the job dispatcher, the agent bridge, the prefix cache and the other users of the bus take holders over the active set. The file manager is built before the set because the NATS file stager needs it. The NATS URL keeps its meaning. A URL the client cannot use still stops the start. A server that is not up does not: the client retries, as it did before. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/agent_jobs.go | 2 +- core/application/application.go | 6 +- core/application/distributed.go | 174 +++++++++++------ core/application/distributed_carrier_test.go | 188 +++++++++++++++++++ core/application/failover_distributed.go | 2 +- core/application/startup.go | 4 +- core/http/app.go | 10 +- core/http/routes/openresponses.go | 2 +- 8 files changed, 315 insertions(+), 73 deletions(-) create mode 100644 core/application/distributed_carrier_test.go diff --git a/core/application/agent_jobs.go b/core/application/agent_jobs.go index f380b075002a..27c1275a7c03 100644 --- a/core/application/agent_jobs.go +++ b/core/application/agent_jobs.go @@ -38,7 +38,7 @@ func (a *Application) RestartAgentJobService() error { agentJobService.SetDistributedJobStore(d.JobStore) } // Keep agent tasks consistent across replicas (same client the dispatcher uses). - agentJobService.SetTaskSyncNATS(d.Nats) + agentJobService.SetTaskSyncNATS(d.Broadcaster) } // Start the service diff --git a/core/application/application.go b/core/application/application.go index 326885b3105b..d22895af6ad7 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -670,8 +670,8 @@ func (a *Application) agentPoolOptions() agentpool.AgentPoolOptions { opts.WorkQueue = d.WorkQueue opts.EventBridge = d.AgentBridge opts.AgentStore = d.AgentStore - if d.Nats != nil { - opts.Bus = d.Nats + if d.Broadcaster != nil { + opts.Bus = d.Broadcaster } } return opts @@ -723,7 +723,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.SetJobSyncNATS(d.Broadcaster) } aps.SetUserServicesManager(usm) diff --git a/core/application/distributed.go b/core/application/distributed.go index 91cb2d053356..bb9b5845e05d 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -7,12 +7,15 @@ import ( "io" "strings" "sync" + "sync/atomic" "time" "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging" @@ -27,8 +30,15 @@ import ( ) // DistributedServices holds all services initialized for distributed mode. +// +// The seams that cross a process boundary are holders over the active carrier +// set: they take the carrier through an atomic pointer, so a later change of +// carrier does not touch the code that holds them. Only Shutdown reaches the +// set itself. type DistributedServices struct { - Nats *messaging.Client + // Broadcaster is the fan-out seam. Code that only publishes and subscribes + // takes this, never the connection of a carrier. + Broadcaster messaging.Broadcaster WorkQueue messaging.WorkQueue AgentControl mcpTools.AgentControl Store storage.ObjectStore @@ -44,13 +54,16 @@ type DistributedServices struct { FileMgr *storage.FileManager FileStager nodes.FileStager ModelAdapter *nodes.ModelRouterAdapter - Unloader *nodes.RemoteUnloaderAdapter + Unloader nodes.NodeControl ModelCleanup *nodes.ModelCleanupService // WorkerHTTPDial reaches a worker's own HTTP server for the admin // backend-logs proxy, the same way the HTTP file stager does. WorkerHTTPDial nodes.WorkerNetDialerFor + // active names the carrier set the holders forward to. + active *atomic.Pointer[carrier.Set] + shutdownOnce sync.Once } @@ -70,10 +83,12 @@ func (ds *DistributedServices) Shutdown() { if closer, ok := ds.Store.(io.Closer); ok { closer.Close() } - // AgentBridge has no Close method — its NATS subscriptions are cleaned up - // when the NATS client is closed below. - if ds.Nats != nil { - ds.Nats.Close() + // AgentBridge has no Close method: its subscriptions are cleaned up + // when the carrier is closed below. + if ds.active != nil { + if set := ds.active.Load(); set != nil && set.Close != nil { + set.Close() + } } xlog.Info("Distributed services shut down") }) @@ -112,7 +127,35 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Distributed instance", "id", cfg.Distributed.InstanceID) - // Connect to NATS + if authDB == nil { + return nil, fmt.Errorf("distributed mode requires auth database to be initialized first") + } + + // The cluster carrier row says which transport the whole cluster uses. + // The first replica to start writes it; later replicas, and every restart, + // follow it. Validate above has already refused a missing NATS URL, so a + // replica that seeds the row seeds NATS. + carrierStore, err := cluster.NewCarrierStore(authDB) + if err != nil { + return nil, fmt.Errorf("initializing cluster carrier state: %w", err) + } + carrierRow, seeded, err := carrierStore.Seed(context.Background(), cluster.CarrierNATS, cfg.Distributed.InstanceID) + if err != nil { + return nil, fmt.Errorf("reading cluster carrier state: %w", err) + } + if seeded { + xlog.Info("Cluster carrier seeded", "carrier", carrierRow.Active, "epoch", carrierRow.Epoch) + } + if carrierRow.Active != cluster.CarrierNATS { + // No silent fallback to the carrier this replica has flags for: a + // cluster that has moved to another carrier would split in two. + return nil, fmt.Errorf("the cluster carrier is %q (epoch %d, set by %s), and this build runs the %q carrier only", + carrierRow.Active, carrierRow.Epoch, carrierRow.ChangedBy, cluster.CarrierNATS) + } + + // Connect to NATS. A URL or credential the client cannot use stops the + // start. A server that is not up yet does not: the client keeps retrying, + // so a frontend can start before its broker. natsAuth := cfg.Distributed.NatsAuthConfig() if natsAuth.RequireAuth && (natsAuth.ServiceUserJWT == "" || natsAuth.ServiceUserSeed == "") { return nil, fmt.Errorf("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") @@ -162,10 +205,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } // Initialize node registry (requires the auth DB which is PostgreSQL) - if authDB == nil { - return nil, fmt.Errorf("distributed mode requires auth database to be initialized first") - } - registry, err := nodes.NewNodeRegistry(authDB) if err != nil { return nil, fmt.Errorf("initializing node registry: %w", err) @@ -207,6 +246,53 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade xlog.Info("Applied declarative model scheduling config", "models", len(schedConfigs)) } + // Initialize file manager with local cache + cacheDir := cfg.DataPath + "/cache" + fileMgr, err := storage.NewFileManager(store, cacheDir) + if err != nil { + return nil, fmt.Errorf("initializing file manager: %w", err) + } + xlog.Info("File manager initialized", "cacheDir", cacheDir) + + // Build the NATS carrier set and put it behind the holders. Everything + // below takes a holder, not the connection. + natsSet, err := carrier.NewNATSSet(carrier.NATSOptions{ + Client: natsClient, + Epoch: carrierRow.Epoch, + Registry: registry, + InstallTimeout: cfg.Distributed.BackendInstallTimeoutOrDefault(), + UpgradeTimeout: cfg.Distributed.BackendUpgradeTimeoutOrDefault(), + Token: cfg.Distributed.RegistrationToken, + S3Staging: cfg.Distributed.StorageURL != "", + FileManager: fileMgr, + HTTPAddrFor: func(nodeID string) (string, error) { + node, err := registry.Get(context.Background(), nodeID) + if err != nil { + return "", err + } + if node.HTTPAddress == "" { + return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) + } + return node.HTTPAddress, nil + }, + }) + if err != nil { + return nil, fmt.Errorf("building the %s carrier: %w", cluster.CarrierNATS, err) + } + active := &atomic.Pointer[carrier.Set]{} + active.Store(natsSet) + broadcaster := carrier.NewBroadcaster(active) + workQueue := carrier.NewWorkQueue(active) + clientFactory := carrier.NewClients(active) + fileStager := carrier.NewFiles(active) + remoteUnloader := carrier.NewCommands(active) + workerHTTPDial := carrier.NewWorkerDialer(active) + if cfg.Distributed.StorageURL != "" { + xlog.Info("File stager initialized (S3+NATS)") + } else { + xlog.Info("File stager initialized (HTTP direct transfer)") + } + // Collect SmartRouter option values; the router itself is created after all // dependencies (including FileStager and Unloader) are ready. var routerAuthToken string @@ -223,6 +309,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade cfg.Distributed.StaleNodeThresholdOrDefault(), routerAuthToken, !cfg.Distributed.DisablePerModelHealthCheck, + clientFactory, ) // Initialize job store @@ -232,10 +319,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Distributed job store initialized") - workQueue := messaging.NewNATSWorkQueue(natsClient) - // Initialize job dispatcher - dispatcher := jobs.NewDispatcher(jobStore, workQueue, natsClient, authDB, cfg.Distributed.InstanceID) + dispatcher := jobs.NewDispatcher(jobStore, workQueue, broadcaster, authDB, cfg.Distributed.InstanceID) // Initialize agent store agentStore, err := agents.NewAgentStore(authDB) @@ -245,7 +330,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade xlog.Info("Distributed agent store initialized") // Initialize agent event bridge - agentBridge := agents.NewEventBridge(natsClient, agentStore, cfg.Distributed.InstanceID) + agentBridge := agents.NewEventBridge(broadcaster, agentStore, cfg.Distributed.InstanceID) // Start observable persister — captures observable_update events from workers // (which have no DB access) and persists them to PostgreSQL. @@ -261,41 +346,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("initializing distributed stores: %w", err) } - // Initialize file manager with local cache - cacheDir := cfg.DataPath + "/cache" - fileMgr, err := storage.NewFileManager(store, cacheDir) - if err != nil { - return nil, fmt.Errorf("initializing file manager: %w", err) - } - xlog.Info("File manager initialized", "cacheDir", cacheDir) - - // Create FileStager for distributed file transfer - workerHTTPDial := nodes.DirectWorkerNetDialer() - var fileStager nodes.FileStager - if cfg.Distributed.StorageURL != "" { - fileStager = nodes.NewS3NATSFileStager(fileMgr, natsClient) - xlog.Info("File stager initialized (S3+NATS)") - } else { - fileStager = nodes.NewHTTPFileStager(func(nodeID string) (string, error) { - node, err := registry.Get(context.Background(), nodeID) - if err != nil { - return "", err - } - if node.HTTPAddress == "" { - return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) - } - return node.HTTPAddress, nil - }, cfg.Distributed.RegistrationToken, workerHTTPDial) - xlog.Info("File stager initialized (HTTP direct transfer)") - } - // Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go - remoteUnloader := nodes.NewRemoteUnloaderAdapter( - registry, - natsClient, - cfg.Distributed.BackendInstallTimeoutOrDefault(), - cfg.Distributed.BackendUpgradeTimeoutOrDefault(), - ) - // Prefix-cache-aware routing. Enabled by default; an operator can opt out // with --distributed-prefix-cache=false, which leaves prefixProvider and // pressure nil so the SmartRouter and reconciler behave exactly as the @@ -315,8 +365,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("invalid prefix-cache configuration: %w", err) } idx := prefixcache.NewIndex(prefixCfg) - prefixSync := prefixcache.NewSync(idx, natsClient) - pressure = prefixcache.NewSyncedPressure(prefixCfg.PressureWindow, natsClient) + prefixSync := prefixcache.NewSync(idx, broadcaster) + pressure = prefixcache.NewSyncedPressure(prefixCfg.PressureWindow, broadcaster) prefixProvider = prefixSync // Invalidate the prefix-cache index whenever a replica row is removed. @@ -342,17 +392,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Apply peers' observations/invalidations to the same Sync. ApplyObserve // and ApplyInvalidate update only the local index and do not re-publish, // so there is no broadcast loop. - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheObserve, func(ev messaging.PrefixCacheObserveEvent) { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCacheObserve, func(ev messaging.PrefixCacheObserveEvent) { prefixSync.ApplyObserve(ev, time.Now()) }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheObserve, err) } - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheInvalidate, func(ev messaging.PrefixCacheInvalidateEvent) { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCacheInvalidate, func(ev messaging.PrefixCacheInvalidateEvent) { prefixSync.ApplyInvalidate(ev) }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err) } - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCachePressure, func(ev messaging.PrefixCachePressureEvent) { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCachePressure, func(ev messaging.PrefixCachePressureEvent) { pressure.ApplyPressure(ev, time.Now()) }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCachePressure, err) @@ -362,7 +412,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // their real KV state without coupling to router internals. Routing stays // on the guessed provider until a backend producer is available. reportedIndex := prefixcache.NewReportedIndex() - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheResidency, reportedIndex.Apply); err != nil { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCacheResidency, reportedIndex.Apply); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheResidency, err) } @@ -407,6 +457,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Unloader: remoteUnloader, ModelCleanup: modelCleanup, FileStager: fileStager, + ClientFactory: clientFactory, GalleriesJSON: routerGalleriesJSON, AuthToken: routerAuthToken, DB: authDB, @@ -454,8 +505,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade "knob", config.FlagDiskHeadroomCheck, "env", "LOCALAI_DISTRIBUTED_DISK_HEADROOM_CHECK") } - router.StagingTracker().SetPublisher(natsClient) - if _, err := router.StagingTracker().SubscribeBroadcasts(natsClient); err != nil { + router.StagingTracker().SetPublisher(broadcaster) + if _, err := router.StagingTracker().SubscribeBroadcasts(broadcaster); err != nil { xlog.Warn("Failed to subscribe to staging progress broadcasts", "error", err) } @@ -468,6 +519,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Scheduler: router, Unloader: remoteUnloader, Adapter: remoteUnloader, + ClientFactory: clientFactory, RegistrationToken: cfg.Distributed.RegistrationToken, DB: authDB, Interval: 30 * time.Second, @@ -483,9 +535,9 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade success = true return &DistributedServices{ - Nats: natsClient, + Broadcaster: broadcaster, WorkQueue: workQueue, - AgentControl: nodes.NewNATSAgentControl(natsClient), + AgentControl: carrier.NewAgents(active), Store: store, Registry: registry, Router: router, @@ -503,6 +555,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade ModelCleanup: modelCleanup, WorkerHTTPDial: workerHTTPDial, + + active: active, }, nil } diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go new file mode 100644 index 000000000000..b5609fb5dbb2 --- /dev/null +++ b/core/application/distributed_carrier_test.go @@ -0,0 +1,188 @@ +package application + +import ( + "context" + "fmt" + "os" + "runtime" + "sync" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" + "gorm.io/gorm" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("distributed startup and the cluster carrier", func() { + var ( + db *gorm.DB + cfg *config.ApplicationConfig + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + cfg = &config.ApplicationConfig{ + Context: ctx, + DataPath: GinkgoT().TempDir(), + Auth: config.AuthConfig{Enabled: true, DatabaseURL: "postgres://unused/db"}, + Distributed: config.DistributedConfig{ + Enabled: true, + InstanceID: "replica-a", + NatsURL: "nats://127.0.0.1:1", + }, + } + }) + + carrierRow := func() cluster.CarrierRow { + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + row, err := store.Get(context.Background()) + Expect(err).ToNot(HaveOccurred()) + return row + } + + It("seeds NATS before it connects, and stops the start on a NATS URL it cannot use", func() { + cfg.Distributed.NatsURL = "nats://[::1" + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("connecting to NATS"))) + Expect(svc).To(BeNil()) + + // The row was written with NATS, because this deployment has a NATS URL. + row := carrierRow() + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + Expect(row.Epoch).To(Equal(int64(1))) + Expect(row.State).To(Equal(cluster.StateStable)) + Expect(row.ChangedBy).To(Equal("replica-a")) + }) + + It("refuses to start on a carrier this build does not run, and does not try NATS", func() { + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(context.Background(), cluster.CarrierTunnel, "replica-b") + Expect(err).ToNot(HaveOccurred()) + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring(`the cluster carrier is "tunnel"`))) + Expect(err.Error()).ToNot(ContainSubstring("connecting to NATS"), "no fallback to the carrier the flags name") + Expect(svc).To(BeNil()) + Expect(carrierRow().Active).To(Equal(cluster.CarrierTunnel), "the row is left as it was") + }) + + It("keeps the existing validation of a missing NATS URL", func() { + cfg.Distributed.NatsURL = "" + _, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("--nats-url"))) + + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, err = store.Get(context.Background()) + Expect(err).To(MatchError(cluster.ErrNotSeeded), "a replica that cannot start leaves no row behind") + }) + + Context("with a NATS server", func() { + BeforeEach(func() { + url, err := startNATS() + if err != nil { + // A CI runner that lost Docker must go red, not skip. Local runs + // without Docker, and macOS CI, skip. + if os.Getenv("CI") != "" && runtime.GOOS != "darwin" { + Fail("testcontainers requires Docker and CI is set: " + err.Error()) + } + Skip("testcontainers requires Docker: " + err.Error()) + } + cfg.Distributed.NatsURL = url + }) + + It("seeds NATS on first start and wires every seam to it", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + Expect(carrierRow().Active).To(Equal(cluster.CarrierNATS)) + Expect(svc.Broadcaster).ToNot(BeNil()) + Expect(svc.WorkQueue).ToNot(BeNil()) + Expect(svc.AgentControl).ToNot(BeNil()) + Expect(svc.FileStager).ToNot(BeNil()) + Expect(svc.Unloader).ToNot(BeNil()) + Expect(svc.WorkerHTTPDial).ToNot(BeNil()) + + // A message published through the holder reaches a subscriber made + // through the same holder, over the real server. + got := make(chan []byte, 1) + sub, err := svc.Broadcaster.Subscribe("nodes.carrier-spec", func(b []byte) { got <- b }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + Expect(svc.Broadcaster.Publish("nodes.carrier-spec", "hello")).To(Succeed()) + Eventually(got).Should(Receive(Equal([]byte(`"hello"`)))) + }) + + It("starts a second replica on the same row without changing it", func() { + first, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(first.Shutdown) + + cfg2 := *cfg + cfg2.Distributed.InstanceID = "replica-b" + cfg2.DataPath = GinkgoT().TempDir() + second, err := initDistributed(&cfg2, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Shutdown) + + row := carrierRow() + Expect(row.Epoch).To(Equal(int64(1))) + Expect(row.ChangedBy).To(Equal("replica-a")) + }) + }) +}) + +var ( + natsOnce sync.Once + natsURL string + natsCtr testcontainers.Container + natsErr error +) + +// startNATS starts one server for the process and returns its URL. +func startNATS() (string, error) { + natsOnce.Do(func() { + ctx := context.Background() + var ctr testcontainers.Container + ctr, natsErr = testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: testcontainers.ContainerRequest{ + Image: "nats:2.10-alpine", + ExposedPorts: []string{"4222/tcp"}, + WaitingFor: wait.ForListeningPort("4222/tcp"), + }, + Started: true, + }) + if natsErr != nil { + return + } + natsCtr = ctr + host, err := ctr.Host(ctx) + if err != nil { + natsErr = err + return + } + port, err := ctr.MappedPort(ctx, "4222/tcp") + if err != nil { + natsErr = err + return + } + natsURL = fmt.Sprintf("nats://%s:%s", host, port.Port()) + }) + return natsURL, natsErr +} + +var _ = AfterSuite(func() { + if natsCtr != nil { + _ = natsCtr.Terminate(context.Background()) + } +}) diff --git a/core/application/failover_distributed.go b/core/application/failover_distributed.go index 87adbdeb51d7..128aa566a8de 100644 --- a/core/application/failover_distributed.go +++ b/core/application/failover_distributed.go @@ -63,7 +63,7 @@ func (a *Application) startFailoverDistributed(ctx context.Context) { xlog.Error("failover: pins will not persist, could not prepare the pin store", "error", err) pins = nil // distsync.New treats a nil store as "no durable pins" } - s, err := distsync.New(ctx, a.distributed.Nats, pins, a.failoverManager) + s, err := distsync.New(ctx, a.distributed.Broadcaster, pins, a.failoverManager) if err != nil { xlog.Error("failover: state will not be shared between frontends", "error", err) return diff --git a/core/application/startup.go b/core/application/startup.go index a9a351dc9906..e77d4e5ba638 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -342,14 +342,14 @@ func New(opts ...config.AppOption) (*Application, error) { 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) + application.agentJobService.SetTaskSyncNATS(distSvc.Broadcaster) } // 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. // Wire NATS and gallery store into GalleryService for cross-instance progress/cancel if application.galleryService != nil { - application.galleryService.SetNATSClient(distSvc.Nats) + application.galleryService.SetNATSClient(distSvc.Broadcaster) if distSvc.DistStores != nil && distSvc.DistStores.Gallery != nil { // Clean up stale in-progress operations from previous crashed instances if _, err := distSvc.DistStores.Gallery.CleanStale(30 * time.Minute); err != nil { diff --git a/core/http/app.go b/core/http/app.go index b6211686baf4..dd51073bda28 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -500,7 +500,7 @@ func API(application *application.Application) (*echo.Echo, error) { // happened to admit, and a load-balanced UI poll alternates between // "operation visible" and "operation gone" between replicas. if d := application.Distributed(); d != nil { - opcache.SetMessagingClient(d.Nats) + opcache.SetMessagingClient(d.Broadcaster) if d.DistStores != nil && d.DistStores.Gallery != nil { opcache.SetGalleryStore(d.DistStores.Gallery) } @@ -518,10 +518,10 @@ func API(application *application.Application) (*echo.Echo, error) { // 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 + var ftNats messaging.Broadcaster var ftStore *distributed.FineTuneStore if d := application.Distributed(); d != nil { - ftNats = d.Nats + ftNats = d.Broadcaster if d.DistStores != nil && d.DistStores.FineTune != nil { ftStore = d.DistStores.FineTune } @@ -540,10 +540,10 @@ func API(application *application.Application) (*echo.Echo, error) { // 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 + var quantNats messaging.Broadcaster var quantStore *distributed.QuantStore if d := application.Distributed(); d != nil { - quantNats = d.Nats + quantNats = d.Broadcaster if d.DistStores != nil && d.DistStores.Quant != nil { quantStore = d.DistStores.Quant } diff --git a/core/http/routes/openresponses.go b/core/http/routes/openresponses.go index 101567a6a683..a137df3b08dc 100644 --- a/core/http/routes/openresponses.go +++ b/core/http/routes/openresponses.go @@ -27,7 +27,7 @@ func RegisterOpenResponsesRoutes(app *echo.Echo, // creator 404s, and the cancel never reaches the CancelFunc (#10993). // Standalone deployments skip this entirely and stay process-local. if err := openresponses.GetGlobalStore().EnableDistributed( - application.ApplicationConfig().Context, d.Nats, application.InstanceID()); err != nil { + application.ApplicationConfig().Context, d.Broadcaster, application.InstanceID()); err != nil { xlog.Error("Failed to enable cross-replica Open Responses store", "error", err) } } From c6b7b336e32199fa70d6c69c0812307bcdfb5469 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:18:54 +0000 Subject: [PATCH 005/162] ci: run the distributed e2e suite as an advisory job make test-e2e-distributed was not run by any workflow. Add one, with continue-on-error so a failure is visible on the pull request and does not block a merge until the job has a track record. The suite starts its own PostgreSQL and NATS servers, so the job pre-pulls both images. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .github/workflows/tests-e2e-distributed.yml | 85 +++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 .github/workflows/tests-e2e-distributed.yml diff --git a/.github/workflows/tests-e2e-distributed.yml b/.github/workflows/tests-e2e-distributed.yml new file mode 100644 index 000000000000..e21dc227b987 --- /dev/null +++ b/.github/workflows/tests-e2e-distributed.yml @@ -0,0 +1,85 @@ +--- +name: 'E2E Distributed Tests' + +on: + pull_request: + # The suite's dependency graph is large, so an allowlist of paths would + # silently stop guarding the moment code moves. It is cheap enough to run + # unless the diff is confined to paths it provably cannot reach. + # See .agents/ci-caching.md. + paths-ignore: + - 'gallery/**' + - 'docs/**' + - 'examples/**' + - '**/*.md' + push: + branches: + - master + +concurrency: + group: ci-tests-e2e-distributed-${{ github.event.pull_request.number || github.sha }}-${{ github.repository }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + tests-e2e-distributed: + runs-on: ubuntu-latest + # Advisory: a failure is a visible red X on the pull request, and the job is + # deliberately not in branch protection. Promoting it to a required check + # is a repository-settings change, to be made once it has a track record. + continue-on-error: true + timeout-minutes: 45 + steps: + - name: Clone + uses: actions/checkout@v7 + with: + submodules: true + - name: Configure apt mirror on runner + uses: ./.github/actions/configure-apt-mirror + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version: '1.26.x' + cache: false + - name: Display Go version + run: go version + - name: Dependencies + run: | + sudo apt-get update + sudo apt-get install -y build-essential libopus-dev + - name: Proto Dependencies + run: | + curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v26.1/protoc-26.1-linux-x86_64.zip -o protoc.zip && \ + unzip -j -d /usr/local/bin protoc.zip bin/protoc && \ + rm protoc.zip + go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2 + go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af + PATH="$PATH:$HOME/go/bin" make protogen-go + - name: Pre-pull test images + # The suite starts its own PostgreSQL and NATS servers with + # testcontainers. Pulling the images here keeps the pull time out of the + # spec timeouts and makes a registry outage read as a setup failure + # instead of a test failure. + run: | + docker pull postgres:16-alpine + docker pull nats:2-alpine + - name: Distributed E2E + # TESTCONTAINERS_RYUK_DISABLED keeps the pre-pull above meaningful. The + # reaper exists to clean up leaked containers on a long-lived host, but + # this runner is ephemeral and every container dies with the VM. + # Leaving it enabled would pull one more unpinned image mid-suite. + env: + TESTCONTAINERS_RYUK_DISABLED: "true" + run: | + PATH="$PATH:$HOME/go/bin" make test-e2e-distributed + # tmate keeps the runner busy until the 6 hour job limit, so a single + # failure costs a whole runner slot. Only open a session when someone + # asked for one by labelling the pull request `ci-debug`, and cap the + # session so a forgotten label cannot idle a runner either. + - name: Setup tmate session if tests fail + if: ${{ failure() && contains(github.event.pull_request.labels.*.name, 'ci-debug') }} + timeout-minutes: 30 + uses: mxschmitt/action-tmate@v3.23 + with: + detached: true + connect-timeout-seconds: 180 + limit-access-to-actor: true From 859f198849cdff7da7ef4f178abc979816910fbc Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:18:54 +0000 Subject: [PATCH 006/162] docs(agents): describe the carrier holders and the cluster row Add a section on the holders, the carrier set, the subscription registry and the cluster_carrier row. Record the import boundary of the NATS libraries and the startup rule for an unreachable NATS server. Remove the open item for a reconciler client factory, which now exists. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 54 ++++++++++++++++++++++++++++++------ 1 file changed, 46 insertions(+), 8 deletions(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 2b313ed1f8eb..cde386d89927 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -3,9 +3,13 @@ Distributed mode talks to other processes through a few interfaces. Code above them must not name NATS, and a new transport is a new implementation of them. Today every seam has one carrier: NATS, or a direct dial for the gRPC and HTTP -connections to a worker. Only the carrier files import `nats.go` -(`core/services/messaging/client.go`, `tls.go` and -`core/services/nodes/control_nats.go`). +connections to a worker. Only the carrier files import the NATS libraries +(`core/services/messaging/client.go`, `tls.go`, +`core/services/nodes/control_nats.go` and `pkg/natsauth`). A spec in +`core/services/carrier` fails when any other non-test file imports +`github.com/nats-io/*`; it reads the import lines of every Go file, so a build +constraint does not hide one. A new carrier adds its library and its files to +that spec. | Seam | Interface | Lives in | Today | |---|---|---|---| @@ -21,6 +25,42 @@ Request and reply payloads of the control verbs live in `core/services/workerctl`, so the frontend half, the worker half and any carrier share one wire format without importing each other. +## Holders and the carrier row + +Code above the seams does not hold a carrier. It holds a forwarding holder from +`core/services/carrier`, one per seam: `Broadcaster`, `WorkQueue`, `Commands` +(`nodes.NodeControl`), `Files`, `Clients`, `NewWorkerDialer` and `Agents`. Each +call loads an `atomic.Pointer[carrier.Set]` once and calls the same method on +that set. A `carrier.Set` is the implementation of every seam for one carrier. +Build a set completely and call `Validate` before you store it. The holders take +no lock and allocate nothing, and a call that is already running is never moved. +`carrier.NewNATSSet` is the one place that builds the NATS set. + +`NodeControl` is the whole control surface a carrier gives the frontend: +`NodeCommandSender`, the process lister, the three unload interfaces the model +loader reads by type assertion, and the install timeout, model file deletion and +forced install that the managers and the reconciler call. Take it, not a +concrete sender. The `Files` holder implements `RequestFileReleaser` because +`FileStagingClient` looks for it, and `carrier.FileCarrier` makes every set +provide it. + +`carrier.Broadcaster` also keeps a registry of every subscription, because a +subscription lives on one carrier. A swap is `Listen(next)` (each handler now +listens on both carriers), then a store of `next` in the pointer (publishes go +to `next`), then `Release(old)`. No message is published on two carriers. +`NotifyReconnect` runs the reconnect hooks that `syncstate` and the failover +sync register, because messages across the flip are not ordered. + +The active carrier is one row of table `cluster_carrier` +(`cluster.CarrierStore`). Each transition is a compare-and-set on its epoch. At +startup a replica reads the row, or inserts it if it is absent, and builds the +set the row names. A replica whose row names a carrier it does not implement +fails to start. It never falls back to another carrier. + +A NATS URL that points at a server which is not up does not stop the start: +`messaging.New` retries on a failed connect, so a frontend can start before its +broker. A URL it cannot parse does. + ## Queues `WorkKind` names the work (`WorkTask`, `WorkMCPCI`, `WorkAgentRun`). @@ -151,10 +191,8 @@ other seams against a real server, also through Docker. - `clientFor` returns no error; a carrier with no dialer for a node needs that path. - The agent worker still uses NATS directly for its connection and for agent - events (`agents.NewEventBridge`). -- `ReplicaReconcilerOptions` has no `ClientFactory` field. Without a - `Prober`, `NewReplicaReconciler` builds its own `tokenClientFactory`, so a - carrier that dials through another factory must add the field. + events. `agents.NewEventBridge` takes a `Broadcaster`, so the frontend's + bridge goes through the holder. - The backend-logs proxy (`proxyHTTPToWorker`) starts from `httpclient.HardenedTransport()`, which keeps `Proxy: http.ProxyFromEnvironment`. With `HTTP_PROXY` set, a dialer that @@ -168,7 +206,7 @@ other seams against a real server, also through Docker. publishes on `events` (`handleMCPCIJob` does this). - Not additive: the agent-run consumer (`NATSDispatcher.runDelivery`) ignores the per-delivery `events` publisher and publishes through the process-wide - `EventBridge`, which is built on the NATS client. A second carrier must + `EventBridge`, which is bound to one `Broadcaster`. A second carrier must change `NATSDispatcher` and `EventBridge`, for example by binding `handleJob`'s publishes to `events` through a bridge view that shares the cancel registry. From 09ea08044ecae393c970a3ad7fb92860e2c822b7 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 16:21:17 +0000 Subject: [PATCH 007/162] test(distributed): run the broadcaster conformance suite on the holder The fan-out holder is a Broadcaster, so it has to pass the suite that every carrier passes. Run it on a holder as built, and on a holder whose carrier was swapped. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../carrier/broadcaster_conformance_test.go | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 core/services/carrier/broadcaster_conformance_test.go diff --git a/core/services/carrier/broadcaster_conformance_test.go b/core/services/carrier/broadcaster_conformance_test.go new file mode 100644 index 000000000000..3d8b0e60345f --- /dev/null +++ b/core/services/carrier/broadcaster_conformance_test.go @@ -0,0 +1,41 @@ +package carrier_test + +import ( + "sync/atomic" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" + . "github.com/onsi/ginkgo/v2" +) + +// The holder is a Broadcaster like any carrier, so it passes the same suite as +// the carriers do: first as it is built, then after its carrier was swapped. +var _ = Describe("Broadcaster holder conformance", func() { + Describe("on the carrier it was built over", func() { + messagingtest.RunBroadcasterConformance(func() (messaging.Broadcaster, func()) { + var cur atomic.Pointer[carrier.Set] + cur.Store(newFakeCarrier(cluster.CarrierNATS, 1).set) + return carrier.NewBroadcaster(&cur), func() {} + }) + }) + + Describe("after the carrier was swapped", func() { + messagingtest.RunBroadcasterConformance(func() (messaging.Broadcaster, func()) { + var cur atomic.Pointer[carrier.Set] + first := newFakeCarrier(cluster.CarrierNATS, 1).set + next := newFakeCarrier(cluster.CarrierTunnel, 2).set + cur.Store(first) + h := carrier.NewBroadcaster(&cur) + if err := h.Listen(next); err != nil { + Fail(err.Error()) + } + cur.Store(next) + if err := h.Release(first); err != nil { + Fail(err.Error()) + } + return h, func() {} + }) + }) +}) From 59c0311157c5c75d2b0124a5e5b603781140eca3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:15:27 +0000 Subject: [PATCH 008/162] feat(distributed): add the cluster registry of replicas and worker connections Each frontend replica records itself in the instances table and refreshes the row from a membership loop. The database clock stamps every row, so no replica decides on the clock of another. The row also carries the version and the epoch of the carrier that the replica has built, with a reason when it could not. The node_connections table records which live replica holds the connection of a worker. Every claim takes a fresh epoch from a sequence, so a replica that lost a worker cannot clear the claim of the replica that won it. A released or orphaned row is kept as a departure, and Presence tells a worker that is reconnecting from one that is gone only after a grace period. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/connected_among_test.go | 171 +++++ core/services/cluster/helpers_test.go | 24 + core/services/cluster/instance.go | 152 ++++ core/services/cluster/instance_test.go | 158 +++++ core/services/cluster/membership.go | 353 ++++++++++ core/services/cluster/membership_test.go | 427 ++++++++++++ core/services/cluster/ownership.go | 454 ++++++++++++ core/services/cluster/ownership_test.go | 650 ++++++++++++++++++ core/services/cluster/presence.go | 171 +++++ core/services/cluster/presence_test.go | 304 ++++++++ 10 files changed, 2864 insertions(+) create mode 100644 core/services/cluster/connected_among_test.go create mode 100644 core/services/cluster/helpers_test.go create mode 100644 core/services/cluster/instance.go create mode 100644 core/services/cluster/instance_test.go create mode 100644 core/services/cluster/membership.go create mode 100644 core/services/cluster/membership_test.go create mode 100644 core/services/cluster/ownership.go create mode 100644 core/services/cluster/ownership_test.go create mode 100644 core/services/cluster/presence.go create mode 100644 core/services/cluster/presence_test.go diff --git a/core/services/cluster/connected_among_test.go b/core/services/cluster/connected_among_test.go new file mode 100644 index 000000000000..5b6e0e4b9e09 --- /dev/null +++ b/core/services/cluster/connected_among_test.go @@ -0,0 +1,171 @@ +package cluster_test + +import ( + "context" + "strings" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// ConnectedAmong is the read a SELECTION is built on: which of these nodes can +// a request reach right now, and which of those can it reach without a relay +// hop. Every spec below is about one of the two ways that answer can be wrong, +// and both are the invariant that the routing code relies on, not tidiness: +// +// - naming a node NO live replica holds sends a request at a process that is +// gone, and the caller reads the failure as the worker's own answer; +// - dropping a node a live replica DOES hold refuses work to a fleet that is +// fine, which is what an absence read leaking into a routing read looks +// like. +var _ = Describe("ConnectedAmong", func() { + var ( + ctx context.Context + db *gorm.DB + reg *cluster.Registry + ) + + // live registers a replica that is heartbeating now. + live := func(id string) { + GinkgoHelper() + Expect(reg.Register(ctx, id, "v1", 0, "")).To(Succeed()) + } + + // kill stops a replica heartbeating, WITHOUT stamping a departure on the + // rows it holds. That is exactly the state a replica that died leaves + // behind until a peer's membership sweep runs. + kill := func(id string) { + GinkgoHelper() + age(ctx, db, "instances", "last_seen", "id", id, cluster.InstanceLiveness+5*time.Second) + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("reports a node this replica holds in BOTH lists", func() { + // heldByOwner is a subset of held and not an alternative to it. A + // caller that prefers what it owns and falls back to the rest would + // otherwise have to union the two itself, and a caller that forgot + // would never fall back at all. + live("me") + _, err := reg.Claim(ctx, "agent-1", "me") + Expect(err).ToNot(HaveOccurred()) + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(ConsistOf("agent-1")) + Expect(byOwner).To(ConsistOf("agent-1")) + }) + + It("reports a node a PEER holds as held, and not as held by this replica", func() { + live("me") + live("peer") + _, err := reg.Claim(ctx, "agent-1", "peer") + Expect(err).ToNot(HaveOccurred()) + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(ConsistOf("agent-1")) + Expect(byOwner).To(BeEmpty()) + }) + + It("drops a node whose owning replica is no longer live", func() { + // The whole reason the statement joins instances. A connection row + // outlives its owner by up to a liveness window plus a heartbeat, so + // without the join this read names a dead replica for that entire + // window and every caller acting on it dials a corpse. + live("dead") + _, err := reg.Claim(ctx, "agent-1", "dead") + Expect(err).ToNot(HaveOccurred()) + kill("dead") + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(BeEmpty()) + Expect(byOwner).To(BeEmpty()) + }) + + It("drops a node whose row records a departure, however recent", func() { + // A departed row is not a candidate for ROUTING, whatever its age. How + // old the departure is decides whether the worker has GONE, which is + // Presence's question and not this one; answering it here would put two + // windows on one fact. + live("me") + epoch, err := reg.Claim(ctx, "agent-1", "me") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "agent-1", "me", epoch)).To(Succeed()) + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(BeEmpty()) + Expect(byOwner).To(BeEmpty()) + }) + + It("answers about the ids it was given and no others", func() { + live("me") + for _, id := range []string{"agent-1", "agent-2", "backend-9"} { + _, err := reg.Claim(ctx, id, "me") + Expect(err).ToNot(HaveOccurred()) + } + + held, byOwner, err := reg.ConnectedAmong(ctx, []string{"agent-1", "agent-2"}, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(ConsistOf("agent-1", "agent-2")) + Expect(byOwner).To(ConsistOf("agent-1", "agent-2")) + }) + + It("answers an empty candidate set without asking the database", func() { + // `IN ()` is a syntax error in PostgreSQL, so a statement issued here + // would fail rather than answer nothing. + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + + held, byOwner, err := recording.ConnectedAmong(ctx, nil, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(held).To(BeEmpty()) + Expect(byOwner).To(BeEmpty()) + Expect(rec.statementCount()).To(Equal(0)) + }) + + It("answers in one joined statement measured on the database clock", func() { + live("me") + _, err := reg.Claim(ctx, "agent-1", "me") + Expect(err).ToNot(HaveOccurred()) + + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + _, _, err = recording.ConnectedAmong(ctx, []string{"agent-1", "agent-2"}, "me") + Expect(err).ToNot(HaveOccurred()) + + sql := strings.ToLower(rec.only()) + // only() rules out the read-per-id shape: between two statements the + // owning replica can die, and the answer would then be assembled from + // two different snapshots of the cluster. + Expect(sql).To(ContainSubstring("join")) + Expect(sql).To(ContainSubstring("instances")) + // The liveness window is computed by the DATABASE. A Go-side cutoff + // would appear here as a bound literal and would then move with each + // replica's clock skew. The test container shares this host's clock, so + // no behavioural spec can catch that; the statement shape is the only + // place it is visible. + Expect(sql).To(ContainSubstring("make_interval")) + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).ToNot(MatchRegexp(`last_seen\s*>\s*'`), + "the liveness cutoff must not be a literal timestamp from this process's clock") + // Held-ness is asked in the SQL rather than left to the join alone. An + // empty owner id matches no instance today only because no replica + // registers under one, which is an accident of who registers rather + // than a property of ownership. + Expect(strings.Join(strings.Fields(sql), " ")).To(ContainSubstring( + "node_connections.owner_instance_id <> ''")) + }) +}) diff --git a/core/services/cluster/helpers_test.go b/core/services/cluster/helpers_test.go new file mode 100644 index 000000000000..32b50f7afebc --- /dev/null +++ b/core/services/cluster/helpers_test.go @@ -0,0 +1,24 @@ +package cluster_test + +import ( + "context" + "errors" + + "github.com/mudler/LocalAI/core/services/cluster" + "gorm.io/gorm" +) + +// ownerRow reads the connection table as it is, without the liveness join that +// Registry.Owner applies. It returns cluster.ErrNoConnection for a node with no +// held row. +func ownerRow(ctx context.Context, db *gorm.DB, nodeID string) (string, int64, error) { + var conn cluster.NodeConnection + err := db.WithContext(ctx).Where("node_id = ? AND owner_instance_id <> ''", nodeID).First(&conn).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return "", 0, cluster.ErrNoConnection + } + if err != nil { + return "", 0, err + } + return conn.OwnerInstanceID, conn.Epoch, nil +} diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go new file mode 100644 index 000000000000..2f5256179e9f --- /dev/null +++ b/core/services/cluster/instance.go @@ -0,0 +1,152 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// ErrInstanceNotFound means no row exists for the requested instance ID. A +// caller uses it to decide between registering again and retrying. +var ErrInstanceNotFound = errors.New("cluster: instance not found") + +// Instance is one live frontend replica, keyed by the ID that the replica chose +// for itself. +type Instance struct { + ID string `gorm:"primaryKey;size:36" json:"id"` + // Development builds include git-describe output and a full commit hash, + // which is longer than 64 bytes after enough commits. + Version string `gorm:"size:255" json:"version"` + LastSeen time.Time `gorm:"index" json:"last_seen"` + // ReadyEpoch is the epoch of the cluster carrier row for which this replica + // has built and checked the carrier it was asked to prepare. It is zero + // until the replica reports one. + ReadyEpoch int64 `gorm:"not null;default:0" json:"ready_epoch"` + // ReadyReason says why the replica is not ready. It is empty when the + // replica is ready. + ReadyReason string `gorm:"size:512;not null;default:''" json:"ready_reason,omitempty"` +} + +// Registry reads and writes the instances table and the table of worker +// connections. +type Registry struct { + db *gorm.DB +} + +// NewRegistry returns a Registry over db. The caller migrates the tables with +// Migrate. +func NewRegistry(db *gorm.DB) *Registry { + return &Registry{db: db} +} + +// Register records this replica and refreshes its LastSeen. It upserts on the +// primary key and does not delete and insert, so a concurrent Live never sees a +// live replica as missing. The readiness columns are written too, so a replica +// that registers again after it was swept does not lose what it reported. +func (r *Registry) Register(ctx context.Context, id, version string, readyEpoch int64, readyReason string) error { + // The database stamps last_seen and not this process. Liveness is compared + // across replicas, so it must be measured on the one clock they share. With + // one clock per replica, the real window becomes the window minus the clock + // error of the writer and the reader, and a healthy peer can look dead. + row := map[string]any{ + "id": id, + "version": version, + "ready_epoch": readyEpoch, + "ready_reason": readyReason, + "last_seen": gorm.Expr("now()"), + } + if err := r.db.WithContext(ctx).Model(&Instance{}).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.Assignments(map[string]any{ + "version": version, + "ready_epoch": readyEpoch, + "ready_reason": readyReason, + "last_seen": gorm.Expr("now()"), + }), + }).Create(row).Error; err != nil { + return fmt.Errorf("registering instance %q: %w", id, err) + } + return nil +} + +// Heartbeat refreshes LastSeen for a registered instance. An unknown ID is an +// error and not an insert, because the caller must then register again with +// the full row. +func (r *Registry) Heartbeat(ctx context.Context, id string) error { + // gorm reports no error when the Where clause matches nothing, so the miss + // is read from RowsAffected. + res := r.db.WithContext(ctx).Model(&Instance{}). + Where("id = ?", id). + Update("last_seen", gorm.Expr("now()")) + if res.Error != nil { + return fmt.Errorf("heartbeating instance %q: %w", id, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("heartbeating instance %q: %w", id, ErrInstanceNotFound) + } + return nil +} + +// ReportReady records that the replica has built the carrier for the cluster +// carrier row at epoch. A non-empty reason says that it could not, and why. +// It also refreshes LastSeen, so a report counts as a heartbeat. +func (r *Registry) ReportReady(ctx context.Context, id string, epoch int64, reason string) error { + res := r.db.WithContext(ctx).Model(&Instance{}). + Where("id = ?", id). + Updates(map[string]any{ + "ready_epoch": epoch, + "ready_reason": reason, + "last_seen": gorm.Expr("now()"), + }) + if res.Error != nil { + return fmt.Errorf("reporting readiness of instance %q: %w", id, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("reporting readiness of instance %q: %w", id, ErrInstanceNotFound) + } + return nil +} + +// instanceIsLive is the one predicate that decides whether a replica is alive. +// Its single bind parameter is the window in seconds. Every reader of that fact +// is written with it: Live lists the rows it selects, Owner refuses an owner +// that it rejects, and ReapStale deletes its negation. Two spellings of one +// fact drift apart, and the drift looks like a relay to a replica that one +// query calls dead and another calls alive. +// +// The column has a table name because Owner reads it across a join, where a +// bare last_seen is ambiguous. +// +// The database computes the cutoff for the same reason it stamps last_seen: a +// reader's own clock must not decide whether another replica is alive. +const instanceIsLive = `instances.last_seen > now() - make_interval(secs => ?)` + +// LiveInstanceIDsSQL selects the ids of the replicas that the deployment +// considers alive. Its single bind parameter is the window in seconds, as for +// instanceIsLive. +// +// It exists for a caller that must decide liveness inside another statement, +// such as the claim queue that releases work held by a replica that is gone and +// must not release work held by a replica that is only slow. If the read and +// the update were two statements, a replica could die or come back between +// them. +// +// It is built from instanceIsLive and does not restate it, so the deployment +// has one spelling of "alive". +const LiveInstanceIDsSQL = `SELECT instances.id FROM instances WHERE ` + instanceIsLive + +// Live returns the instances whose LastSeen is newer than now minus within. +func (r *Registry) Live(ctx context.Context, within time.Duration) ([]Instance, error) { + var out []Instance + if err := r.db.WithContext(ctx). + Where(instanceIsLive, within.Seconds()). + Order("id"). + Find(&out).Error; err != nil { + return nil, fmt.Errorf("listing live instances: %w", err) + } + return out, nil +} diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go new file mode 100644 index 000000000000..21944c14a8e8 --- /dev/null +++ b/core/services/cluster/instance_test.go @@ -0,0 +1,158 @@ +package cluster_test + +import ( + "context" + "strings" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("Instance registry", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + // get reads one row as it is stored. + get := func(id string) cluster.Instance { + GinkgoHelper() + var got cluster.Instance + Expect(db.WithContext(ctx).Where("id = ?", id).First(&got).Error).To(Succeed()) + return got + } + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("registers an instance and reads it back", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].ID).To(Equal("inst-a")) + Expect(live[0].Version).To(Equal("v1")) + }) + + It("preserves long development build versions", func() { + version := "v4.10.0-183-g46c57bf3d (" + strings.Repeat("a", 40) + ")" + Expect(len(version)).To(BeNumerically(">", 64)) + Expect(reg.Register(ctx, "inst-dev", version, 0, "")).To(Succeed()) + + Expect(get("inst-dev").Version).To(Equal(version)) + }) + + It("registering the same id again updates the row and does not duplicate it", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "inst-a", "v2", 0, "")).To(Succeed()) + + live, err := reg.Live(ctx, time.Hour) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].Version).To(Equal("v2")) + }) + + It("excludes instances whose heartbeat has aged out", func() { + Expect(reg.Register(ctx, "stale", "v1", 0, "")).To(Succeed()) + // Age the row directly. A spec does not sleep. + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "stale"). + Update("last_seen", gorm.Expr("now() - interval '10 minutes'")).Error).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(BeEmpty()) + }) + + It("lists only the live replicas in LiveInstanceIDsSQL, with the same window", func() { + Expect(reg.Register(ctx, "live", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "stale", "v1", 0, "")).To(Succeed()) + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "stale"). + Update("last_seen", gorm.Expr("now() - interval '10 minutes'")).Error).To(Succeed()) + + var ids []string + Expect(db.Raw(cluster.LiveInstanceIDsSQL, time.Minute.Seconds()).Scan(&ids).Error).To(Succeed()) + Expect(ids).To(Equal([]string{"live"})) + }) + + It("brings a stale instance back with a heartbeat", func() { + Expect(reg.Register(ctx, "revive", "v1", 0, "")).To(Succeed()) + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "revive"). + Update("last_seen", gorm.Expr("now() - interval '10 minutes'")).Error).To(Succeed()) + Expect(reg.Heartbeat(ctx, "revive")).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + }) + + It("heartbeating an unknown instance is an error and not a silent insert", func() { + Expect(reg.Heartbeat(ctx, "ghost")).To(MatchError(cluster.ErrInstanceNotFound)) + live, err := reg.Live(ctx, time.Hour) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(BeEmpty()) + }) + + Describe("readiness for a carrier change", func() { + It("starts at epoch zero with no reason", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + + got := get("inst-a") + Expect(got.ReadyEpoch).To(BeZero()) + Expect(got.ReadyReason).To(BeEmpty()) + }) + + It("records the epoch a replica is ready for", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + + Expect(reg.ReportReady(ctx, "inst-a", 7, "")).To(Succeed()) + + got := get("inst-a") + Expect(got.ReadyEpoch).To(Equal(int64(7))) + Expect(got.ReadyReason).To(BeEmpty()) + }) + + It("records why a replica could not build the carrier", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + + Expect(reg.ReportReady(ctx, "inst-a", 7, "cannot listen on the database")).To(Succeed()) + + Expect(get("inst-a").ReadyReason).To(Equal("cannot listen on the database")) + }) + + It("counts a report as a heartbeat", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "inst-a"). + Update("last_seen", gorm.Expr("now() - interval '10 minutes'")).Error).To(Succeed()) + + Expect(reg.ReportReady(ctx, "inst-a", 2, "")).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + }) + + It("refuses a report from a replica that is not registered", func() { + Expect(reg.ReportReady(ctx, "ghost", 2, "")).To(MatchError(cluster.ErrInstanceNotFound)) + }) + + It("is not changed by a heartbeat", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + Expect(reg.ReportReady(ctx, "inst-a", 5, "")).To(Succeed()) + + Expect(reg.Heartbeat(ctx, "inst-a")).To(Succeed()) + + Expect(get("inst-a").ReadyEpoch).To(Equal(int64(5))) + }) + }) +}) diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go new file mode 100644 index 000000000000..cbec21219f3b --- /dev/null +++ b/core/services/cluster/membership.go @@ -0,0 +1,353 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "math" + "sync" + "time" + + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +const ( + // InstanceHeartbeat is how often a replica refreshes its own row. + InstanceHeartbeat = 5 * time.Second + + // InstanceLiveness is how long a replica may go without a heartbeat before + // its peers treat it as gone: six heartbeats. + // + // The window is long on purpose. When a replica is declared dead, a + // departure is recorded for every connection row it owned, and a worker + // that is recorded as departed while its owner is only slow must be + // re-homed for nothing. Waiting costs little: traffic for that worker is + // retried and not lost. + InstanceLiveness = 30 * time.Second + + // deregisterTimeout bounds the deregistration that Stop does. Shutdown must + // not wait for a database for long. + deregisterTimeout = 5 * time.Second + + // DepartedRetention is the lowest time that a connection row outlives the + // tunnel it recorded before the sweep deletes it. DepartedRetentionFor is + // what the sweep uses. + // + // It is a multiple of the liveness window and not the window itself. The + // row stays so that the age of a departure can be answered at all. If the + // retention were as short as the window that a reader compares the age + // with, the sweep could delete the row under the reader, and a worker that + // is dialling again would look like a worker that was never here. + DepartedRetention = 10 * InstanceLiveness + + // departedRetentionGraceFactor is how many reconnect graces a departure is + // kept for when the grace is the larger of the two values. + departedRetentionGraceFactor = 5 +) + +// DepartedRetentionFor returns how long a departure must be kept, for the +// reconnect grace that the readers of the departure compare it with. +// +// The retention is derived and not fixed because two different people set the +// two windows. If an operator raises the grace above a fixed retention, the +// sweep deletes departures before the grace has passed. Presence then answers +// unknown for that worker for ever, and nothing reaps its rows. +func DepartedRetentionFor(grace time.Duration) time.Duration { + // Check before the multiplication. A grace of more than about 58 years + // overflows the nanoseconds of a Duration, and the product is then negative + // or small, which brings back the defect that this function removes. + if grace > time.Duration(math.MaxInt64)/departedRetentionGraceFactor { + return time.Duration(math.MaxInt64) + } + if scaled := departedRetentionGraceFactor * grace; scaled > DepartedRetention { + return scaled + } + return DepartedRetention +} + +// Membership publishes this replica in the instances table and removes the +// replicas that have stopped answering. +// +// It is the only writer of the row of this replica and the only sweeper of the +// rows of the others. This keeps one fact on one clock: whether a replica is +// alive is answered by its last_seen and by nothing else. +type Membership struct { + reg *Registry + id string + version string + + interval time.Duration + liveness time.Duration + + stop chan struct{} + done chan struct{} + stopOnce sync.Once + + // mu guards the fields below. started tells Stop whether there is a loop to + // join. retention is how long a departure is kept before the sweep deletes + // it, derived from the reconnect grace so that the sweep never outruns it. + // The ready fields are what this replica last reported, so that a + // registration after a sweep writes them back. + mu sync.Mutex + started bool + retention time.Duration + readyEpoch int64 + readyReason string +} + +// NewMembership returns the membership loop for one replica. +func NewMembership(reg *Registry, id, version string) *Membership { + return &Membership{ + reg: reg, + id: id, + version: version, + interval: InstanceHeartbeat, + liveness: InstanceLiveness, + retention: DepartedRetention, + stop: make(chan struct{}), + done: make(chan struct{}), + } +} + +// SetReconnectGrace tells the loop the window that Presence compares a +// departure with, so that the sweep keeps departures for longer than any reader +// needs them. Without it the sweep uses DepartedRetention, which is correct for +// every grace up to that value. +// +// It is safe on a nil receiver, like Stop. +func (m *Membership) SetReconnectGrace(grace time.Duration) { + if m == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + m.retention = DepartedRetentionFor(grace) +} + +// ReportReady records in the instances table that this replica has built the +// carrier for the cluster carrier row at epoch, or, with a reason, that it +// could not. The loop writes the same values again whenever it has to register +// the replica again. The replica must be started. +func (m *Membership) ReportReady(ctx context.Context, epoch int64, reason string) error { + if err := m.reg.ReportReady(ctx, m.id, epoch, reason); err != nil { + return err + } + m.mu.Lock() + m.readyEpoch, m.readyReason = epoch, reason + m.mu.Unlock() + return nil +} + +// register writes the row of this replica with what it last reported. +func (m *Membership) register(ctx context.Context) error { + m.mu.Lock() + epoch, reason := m.readyEpoch, m.readyReason + m.mu.Unlock() + return m.reg.Register(ctx, m.id, m.version, epoch, reason) +} + +// Start registers this replica, then heartbeats and sweeps in the background. +// The first registration is synchronous and its failure is returned: a replica +// that never reaches the table is invisible to its peers, and a background log +// line would hide that. +func (m *Membership) Start(ctx context.Context) error { + if err := m.register(ctx); err != nil { + return err + } + xlog.Info("Cluster instance registered", "id", m.id) + m.mu.Lock() + m.started = true + m.mu.Unlock() + go m.loop(ctx) + return nil +} + +// Stop ends the loop, waits for it, and removes the row of this replica. +// +// Deregistering makes a rolling restart quick for the others. A replica that +// only closes its sockets looks like one that crashed, and its peers go on +// dialling it for the whole liveness window. It cannot be more than best effort, +// because a killed process never gets here, so the sweep still exists. +// +// It is safe to call more than once, and on a Membership that was never started. +func (m *Membership) Stop() { + if m == nil { + return + } + m.mu.Lock() + started := m.started + m.mu.Unlock() + if started { + m.stopOnce.Do(func() { close(m.stop) }) + // Only a started Membership closes done. Waiting on one that never + // started would block for ever. + <-m.done + } + + // This is not the context that Start received: that is the context of the + // application, and it is usually cancelled by the time Stop runs. + ctx, cancel := context.WithTimeout(context.Background(), deregisterTimeout) + defer cancel() + if err := m.reg.Deregister(ctx, m.id); err != nil { + xlog.Warn("Deregistering this replica failed; peers drop it when its heartbeat ages out", + "id", m.id, "within", m.liveness, "error", err) + return + } + xlog.Info("Cluster instance deregistered", "id", m.id) +} + +func (m *Membership) loop(ctx context.Context) { + defer close(m.done) + + ticker := time.NewTicker(m.interval) + defer ticker.Stop() + + for { + select { + case <-m.stop: + return + case <-ctx.Done(): + return + case <-ticker.C: + m.tick(ctx) + } + } +} + +// tick refreshes the row of this replica and sweeps the dead ones. +// +// Every replica sweeps and there is no elected sweeper. The deletes are +// idempotent and cheap, and an elected sweeper is one more process that must be +// alive for the cluster to notice that another is not. +func (m *Membership) tick(ctx context.Context) { + err := m.reg.Heartbeat(ctx, m.id) + if errors.Is(err, ErrInstanceNotFound) { + // Another replica swept this row while this process stalled for long + // enough to look dead. Register again and do not only heartbeat: the + // row has to be rebuilt. + xlog.Warn("Cluster instance row was reaped, registering again", "id", m.id) + if err := m.register(ctx); err != nil { + xlog.Error("Registering the cluster instance again failed", "id", m.id, "error", err) + } + } else if err != nil { + xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) + } + + instances, cleared, err := m.reg.ReapStale(ctx, m.id, m.liveness) + if err != nil { + xlog.Warn("Reaping stale cluster instances failed", "error", err) + return + } + if instances > 0 || cleared > 0 { + // Two counts because the sweep does two things: the instance rows are + // deleted, and the connection rows stay and now record a departure. + xlog.Info("Swept cluster state left by dead replicas", + "instances_deleted", instances, "connections_departed", cleared) + } + + // The sweep also owns the retention. A departure that nothing deletes is a + // row for every worker that ever dialled this deployment. + m.mu.Lock() + retention := m.retention + m.mu.Unlock() + purged, err := m.reg.PurgeDepartedBefore(ctx, retention) + if err != nil { + xlog.Warn("Purging departed worker connections failed", "error", err) + return + } + if purged > 0 { + xlog.Info("Purged worker connections whose departure aged out", "connections", purged, "retention", retention) + } +} + +// Deregister removes one replica and records a departure for every connection +// that it owned. +// +// Both happen in one transaction, as in ReapStale. A replica that is gone owns +// nothing, and a connection row that still names it would send every reader to +// an owner that does not exist. This is the announced form of what the sweep +// does by inference, and the two must agree on what "gone" leaves behind. +// +// The connection rows are cleared and not deleted. A worker whose frontend shut +// down is about to dial the load balancer again, and without its row the +// seconds in between look like a worker that was never connected. +// +// The instances are changed first and the connections second. ReapStale takes +// the same order. The two run at the same time in the usual case, a replica +// shutting down while a peer sweeps it, and both lock the same two tables. With +// opposite orders each could wait for a row held by the other. +func (r *Registry) Deregister(ctx context.Context, id string) error { + if err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + // A row that another replica already swept is the normal result of a slow + // shutdown and not an error, so RowsAffected is not checked. + if err := tx.Where("id = ?", id).Delete(&Instance{}).Error; err != nil { + return fmt.Errorf("deleting instance %q: %w", id, err) + } + // Held rows only. An empty id would match every departed row in the table + // and reset the age of every departure. + if err := tx.Model(&NodeConnection{}). + Where("owner_instance_id = ? AND "+connectionIsHeld, id). + Updates(departure()).Error; err != nil { + return fmt.Errorf("recording departures for connections owned by %q: %w", id, err) + } + return nil + }); err != nil { + return fmt.Errorf("deregistering instance %q: %w", id, err) + } + return nil +} + +// ReapStale deletes the replicas that have not heartbeated within the liveness +// window, and records a departure for every connection row whose owner is no +// longer among the survivors. +// +// It is one sweeper for both. A connection row is orphaned only when its owner +// dies, so the moment to clean up after it is the moment the death is decided. +// A second sweeper with its own schedule would lag behind this one or race it, +// and would need its own answer to "is that replica alive". +// +// The connection rows are cleared and not deleted, which is why the second +// result counts rows that were cleared. A worker whose owner died is dialling +// the load balancer again, and deleting its row would erase the departure that +// says how long ago that began. +// +// self is never reaped. This process can fail to heartbeat for longer than the +// window (a long stall, a database error) and still serve. Deleting its own row +// would mark as departed the workers that are connected to it at that moment. +// The protection goes one way: another replica can still reap a replica that +// stalls for long enough. +// +// It needs PostgreSQL, like Live, and the interval is measured on the clock of +// the database. +func (r *Registry) ReapStale(ctx context.Context, self string, within time.Duration) (instances int64, cleared int64, err error) { + err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + // "Stale" is the negation of the live predicate and not a second + // comparison, so both treat a row with a NULL last_seen the same way. + res := tx.Where("id <> ? AND NOT ("+instanceIsLive+")", self, within.Seconds()). + Delete(&Instance{}) + if res.Error != nil { + return fmt.Errorf("deleting stale instances: %w", res.Error) + } + instances = res.RowsAffected + + // What survived the delete above is the live set. This needs no second + // liveness rule, so it cannot disagree with the first. + // + // Held rows only. An empty owner is in no instance id, so a departed row + // also matches the set difference, and clearing it again at every sweep + // would move its departure forward and keep it from ever ageing out. + res = tx.Model(&NodeConnection{}). + Where("owner_instance_id NOT IN (SELECT id FROM instances) AND " + connectionIsHeld). + Updates(departure()) + if res.Error != nil { + return fmt.Errorf("recording departures for orphaned node connections: %w", res.Error) + } + cleared = res.RowsAffected + return nil + }) + if err != nil { + return 0, 0, fmt.Errorf("reaping stale cluster state: %w", err) + } + return instances, cleared, nil +} diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go new file mode 100644 index 000000000000..ee1c90b7e84d --- /dev/null +++ b/core/services/cluster/membership_test.go @@ -0,0 +1,427 @@ +package cluster_test + +import ( + "context" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("Reaping dead replicas", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + // age pushes a replica's heartbeat into the past. Sleeping in a spec is + // forbidden, and the liveness window is measured in tens of seconds. + age := func(id string, by time.Duration) { + GinkgoHelper() + Expect(db.Model(&cluster.Instance{}).Where("id = ?", id). + Update("last_seen", gorm.Expr("now() - make_interval(secs => ?)", by.Seconds())).Error).To(Succeed()) + } + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("deletes a replica that stopped heartbeating, and the connections it owned", func() { + Expect(reg.Register(ctx, "live", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "dead", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "dead") + Expect(err).ToNot(HaveOccurred()) + age("dead", time.Hour) + + instances, connections, err := reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(instances).To(Equal(int64(1))) + Expect(connections).To(Equal(int64(1)), + "a worker whose owner no longer exists is recorded as connected to nothing") + + _, _, err = ownerRow(ctx, db, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("records a departure rather than erasing the connection when a stale replica is swept", func() { + // The row is what tells "this worker's owner died a moment ago" from + // "this worker has never connected". Deleting it on the sweep erased + // exactly the departure a reconnect grace has to be measured from. + Expect(reg.Register(ctx, "live", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "dead", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "dead") + Expect(err).ToNot(HaveOccurred()) + age("dead", time.Hour) + + _, _, err = reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + + var row cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&row).Error).To(Succeed(), + "the sweep deleted the row, so the worker its owner left behind looks like one that never dialled") + Expect(row.OwnerInstanceID).To(BeEmpty()) + Expect(row.DisconnectedAt).ToNot(BeNil()) + }) + + It("does not restamp a departure it has already recorded", func() { + // The sweep runs every heartbeat. Re-clearing a row it already cleared + // would push its departure forward on every pass, so the departure + // would never age out of any window and the row would be reported as + // swept for as long as it existed. + Expect(reg.Register(ctx, "live", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "dead", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "dead") + Expect(err).ToNot(HaveOccurred()) + age("dead", time.Hour) + + _, connections, err := reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(connections).To(Equal(int64(1))) + var first cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&first).Error).To(Succeed()) + + _, connections, err = reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(connections).To(BeZero(), "the sweeper reported clearing a connection that was already departed") + + var second cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&second).Error).To(Succeed()) + Expect(second.DisconnectedAt).ToNot(BeNil()) + Expect(*second.DisconnectedAt).To(Equal(*first.DisconnectedAt), + "the second sweep moved the departure forward, so it can never age past a grace window") + }) + + It("records a departure rather than erasing the connection when a replica deregisters", func() { + // The announced form of the same thing the sweeper does by inference, + // and the two must not disagree about what "gone" leaves behind. + Expect(reg.Register(ctx, "leaving", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "leaving") + Expect(err).ToNot(HaveOccurred()) + + Expect(reg.Deregister(ctx, "leaving")).To(Succeed()) + + var row cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&row).Error).To(Succeed()) + Expect(row.OwnerInstanceID).To(BeEmpty()) + Expect(row.DisconnectedAt).ToNot(BeNil()) + }) + + It("leaves the connections of a live replica alone", func() { + Expect(reg.Register(ctx, "live", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "other", "v1", 0, "")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "other") + Expect(err).ToNot(HaveOccurred()) + + _, connections, err := reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(connections).To(BeZero()) + + owner, stored, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("other")) + Expect(stored).To(Equal(epoch)) + }) + + It("never reaps the sweeper itself, however stale its own row looks", func() { + // A replica whose heartbeat stalled longer than the window is still + // serving the workers connected to it. Reaping its own row would delete + // their connection rows in the same pass, re-homing workers that never + // went anywhere. + Expect(reg.Register(ctx, "me", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "me") + Expect(err).ToNot(HaveOccurred()) + age("me", time.Hour) + + instances, connections, err := reg.ReapStale(ctx, "me", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(instances).To(BeZero()) + Expect(connections).To(BeZero()) + + owner, _, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me")) + }) + + It("deregisters a replica and the connections it owned, so peers drop it at once", func() { + // Without this a cleanly stopped replica is indistinguishable from a + // crashed one, and every peer keeps dialling it for the whole liveness + // window. + Expect(reg.Register(ctx, "leaving", "v1", 0, "")).To(Succeed()) + Expect(reg.Register(ctx, "staying", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "leaving") + Expect(err).ToNot(HaveOccurred()) + _, err = reg.Claim(ctx, "w2", "staying") + Expect(err).ToNot(HaveOccurred()) + + Expect(reg.Deregister(ctx, "leaving")).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].ID).To(Equal("staying")) + + // The same rule the sweeper applies: a replica that is gone owns + // nothing, and a claim naming it would point every reader at an owner + // that no longer exists. + _, _, err = ownerRow(ctx, db, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + owner, _, err := ownerRow(ctx, db, "w2") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("staying"), "deregistering one replica took another replica's claim") + }) + + It("does not restamp a departure when a replica with no id deregisters", func() { + // Deregister matches an owner id, and a departed row carries an empty + // one, so an empty id would match every departure in the table and reset + // each one's age. Nothing generates an empty id today, which is exactly + // why the filter has to be in the query rather than in that habit. + Expect(reg.Register(ctx, "leaving", "v1", 0, "")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "leaving") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "leaving", epoch)).To(Succeed()) + var before cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&before).Error).To(Succeed()) + + Expect(reg.Deregister(ctx, "")).To(Succeed()) + + var after cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&after).Error).To(Succeed()) + Expect(after.DisconnectedAt).ToNot(BeNil()) + Expect(*after.DisconnectedAt).To(Equal(*before.DisconnectedAt), + "a departure that had already been recorded was stamped again, so its age restarted") + }) + + It("purges a departure that has aged past the retention, from the loop that sweeps it", func() { + // The retention exists only if something applies it, and this loop is + // the only thing that does. A PurgeDepartedBefore nobody calls leaves a + // row per worker that ever dialled this deployment, which is the state + // the sweep's held-ness filter exists to make reachable in the first + // place. + Expect(reg.Register(ctx, "me", "v1", 0, "")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "me") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "me", epoch)).To(Succeed()) + // Aged on the database clock, past the retention the loop applies. + Expect(db.WithContext(ctx).Exec( + `UPDATE node_connections SET disconnected_at = now() - make_interval(secs => ?) WHERE node_id = ?`, + (cluster.DepartedRetention + time.Minute).Seconds(), "w1").Error).To(Succeed()) + + membership := cluster.NewMembership(reg, "me", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + // Polled rather than slept: the loop ticks on its own schedule, and the + // spec is about the call happening at all, not about when. + Eventually(func() (int64, error) { + var rows int64 + err := db.WithContext(ctx).Model(&cluster.NodeConnection{}).Count(&rows).Error + return rows, err + }, "20s", "250ms").Should(BeZero(), + "the loop that sweeps dead replicas never applied the departure retention") + }) + + It("purges on the retention its reconnect grace requires, not on the floor", func() { + // The wiring, not the arithmetic. DepartedRetentionFor is pinned on its + // own in presence_test.go, and a sweep that ignored it and passed the + // DepartedRetention constant would satisfy every one of those + // expectations: the derived value is only worth anything if the loop + // that deletes rows is the thing reading it. + // + // A grace of 10m derives a retention of 50m, well past the 5m floor, so + // the two rows below straddle the difference and one answer cannot + // stand in for the other. + const grace = 10 * time.Minute + Expect(reg.Register(ctx, "me", "v1", 0, "")).To(Succeed()) + depart := func(node string, ago time.Duration) { + epoch, err := reg.Claim(ctx, node, "me") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, node, "me", epoch)).To(Succeed()) + // Aged on the database clock, like every window in this package. + Expect(db.WithContext(ctx).Exec( + `UPDATE node_connections SET disconnected_at = now() - make_interval(secs => ?) WHERE node_id = ?`, + ago.Seconds(), node).Error).To(Succeed()) + } + // Past the 5m floor and inside the 50m the grace derives. A sweep + // applying the floor deletes this row, which is the whole defect: the + // purge would be outrunning the window Presence measures against, and + // a worker that really is gone would read as unknown for ever after. + depart("w-inside-derived-retention", cluster.DepartedRetention+time.Minute) + // Past both, so its deletion is the WITNESS that the sweep ran at all. + // Without it "the row survived" would also be true of a loop that never + // ticked, and the spec would pass on nothing happening. + depart("w-past-derived-retention", cluster.DepartedRetentionFor(grace)+time.Minute) + + membership := cluster.NewMembership(reg, "me", "v1") + membership.SetReconnectGrace(grace) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + // One assertion for both halves: rows are only ever deleted, so the set + // settles, and the settled set says which retention the loop applied. + Eventually(func() ([]string, error) { + var ids []string + err := db.WithContext(ctx).Model(&cluster.NodeConnection{}). + Order("node_id").Pluck("node_id", &ids).Error + return ids, err + }, "20s", "250ms").Should(Equal([]string{"w-inside-derived-retention"}), + "the sweep applied a retention that is not the one this replica's reconnect grace derives") + }) + + It("deregisters when the membership loop stops", func() { + membership := cluster.NewMembership(reg, "me", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + + membership.Stop() + + live, err = reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(BeEmpty(), "a replica that shut down cleanly left its row behind for peers to dial") + }) + + It("takes the two tables in one order, shared with the sweeper, so the two cannot deadlock", func() { + // A replica deregistering and a peer sweeping it run concurrently by + // design, and both lock rows in instances and in node_connections. In + // opposite orders each can end up holding the row the other waits for. + // The order is asserted on the SQL because the alternative, racing two + // transactions until they actually deadlock, is exactly the flaky spec + // this one replaces. + Expect(reg.Register(ctx, "leaving", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "leaving") + Expect(err).ToNot(HaveOccurred()) + + deregRec := newSQLRecorder() + Expect(cluster.NewRegistry(db.Session(&gorm.Session{Logger: deregRec})). + Deregister(ctx, "leaving")).To(Succeed()) + + reapRec := newSQLRecorder() + _, _, err = cluster.NewRegistry(db.Session(&gorm.Session{Logger: reapRec})). + ReapStale(ctx, "sweeper", time.Minute) + Expect(err).ToNot(HaveOccurred()) + + Expect(deregRec.writeOrder()).To(Equal([]string{"instances", "node_connections"})) + Expect(reapRec.writeOrder()).To(Equal(deregRec.writeOrder()), + "the sweeper and deregistration must lock the same two tables in the same order") + }) + + It("tolerates a repeated deregistration, because a sweeper may have got there first", func() { + Expect(reg.Register(ctx, "gone", "v1", 0, "")).To(Succeed()) + Expect(reg.Deregister(ctx, "gone")).To(Succeed()) + Expect(reg.Deregister(ctx, "gone")).To(Succeed()) + }) + + It("stops safely when it was never started", func() { + // The loop channel is only closed by a started loop, so joining an + // unstarted one would block for ever. + membership := cluster.NewMembership(reg, "never-started", "v1") + done := make(chan struct{}) + go func() { + defer GinkgoRecover() + defer close(done) + membership.Stop() + }() + Eventually(done, "10s").Should(BeClosed()) + }) + + It("keeps this replica's row alive and reaps the dead while it runs", func() { + Expect(reg.Register(ctx, "dead", "v1", 0, "")).To(Succeed()) + age("dead", time.Hour) + + membership := cluster.NewMembership(reg, "me", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + // Rows, not live rows: an aged-out replica drops out of Live + // immediately, and what the sweeper adds is deleting it. Asserting on + // Live here would pass with no sweeper at all. + rows := func() int64 { + var n int64 + if err := db.Model(&cluster.Instance{}).Count(&n).Error; err != nil { + return -1 + } + return n + } + Expect(rows()).To(Equal(int64(2)), "the stale row is still in the table until a sweep deletes it") + + Eventually(rows, 3*cluster.InstanceHeartbeat, time.Second).Should(Equal(int64(1))) + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].ID).To(Equal("me"), "the sweeper deleted the wrong row") + }) +}) + +var _ = Describe("Reporting readiness through the membership loop", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + readiness := func(id string) (int64, string) { + GinkgoHelper() + var got cluster.Instance + Expect(db.WithContext(ctx).Where("id = ?", id).First(&got).Error).To(Succeed()) + return got.ReadyEpoch, got.ReadyReason + } + + It("writes what the replica reported", func() { + membership := cluster.NewMembership(reg, "me", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + Expect(membership.ReportReady(ctx, 4, "")).To(Succeed()) + + epoch, reason := readiness("me") + Expect(epoch).To(Equal(int64(4))) + Expect(reason).To(BeEmpty()) + }) + + It("writes it again when the loop has to register the replica after a sweep", func() { + membership := cluster.NewMembership(reg, "me", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + Expect(membership.ReportReady(ctx, 4, "slow disk")).To(Succeed()) + // A peer swept the row while this process stalled. + Expect(reg.Deregister(ctx, "me")).To(Succeed()) + + Eventually(func() int64 { + var n int64 + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "me").Count(&n).Error).To(Succeed()) + return n + }, 3*cluster.InstanceHeartbeat, 250*time.Millisecond).Should(Equal(int64(1))) + + epoch, reason := readiness("me") + Expect(epoch).To(Equal(int64(4)), "the replica lost the epoch it had reported") + Expect(reason).To(Equal("slow disk")) + }) + + It("does not remember a report that the database refused", func() { + membership := cluster.NewMembership(reg, "never-registered", "v1") + + Expect(membership.ReportReady(ctx, 4, "")).To(MatchError(cluster.ErrInstanceNotFound)) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + epoch, _ := readiness("never-registered") + Expect(epoch).To(BeZero()) + }) +}) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go new file mode 100644 index 000000000000..b4040c83e093 --- /dev/null +++ b/core/services/cluster/ownership.go @@ -0,0 +1,454 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// ErrNoConnection reports that no tunnel is recorded for the requested node, or +// that the claim a caller named is no longer the live one. Callers distinguish +// it from a transport failure to decide whether to relay through an owner, to +// answer "this worker is not connected here", or to retry. +var ErrNoConnection = errors.New("cluster: no connection recorded for node") + +// isPostgres reports whether the gorm dialect is PostgreSQL. The connection +// fence is built out of PostgreSQL-only pieces (a sequence, ON CONFLICT +// RETURNING), and the single-binary path runs on SQLite. advisorylock has an +// identical private check; this one is copied rather than imported because it +// is a one-line comparison on gorm's own dialector name and cluster is +// deliberately a leaf package. +func isPostgres(db *gorm.DB) bool { + return strings.Contains(db.Dialector.Name(), "postgres") +} + +// epochSequence is the PostgreSQL sequence every claim draws its epoch from. +// A sequence rather than a per-row counter because a connection row does not +// outlive the deployment: a departure is purged once it is old enough, and with +// `epoch = epoch + 1` the numbering would restart at 1 for the next claim, so a +// replica that claimed, lost the worker, and claimed again could be handed an +// epoch it already held, and a delayed cleanup from the first claim would then +// match, and clear, the live one. +const epochSequence = "node_connection_epochs" + +// NodeConnection records which frontend replica currently holds a worker's +// tunnel. There is at most one row per node: a worker holds exactly one link, +// and whoever wrote the row last owns it. +// +// Epoch is the fence. A worker whose link is silently broken reconnects and may +// land on another replica before the previous owner's socket has noticed, so +// for a while two replicas both believe they own it. Every claim draws a fresh, +// never-reused epoch, so the loser can be told apart from the winner by a number +// both of them hold, without either having to detect the broken socket first. +// +// There is deliberately no last-seen column here. Whether the owning replica is +// alive is answered by Instance.LastSeen, and whether a claim is still the live +// one is answered by the epoch; a second liveness clock for the same fact would +// only drift from the first. +type NodeConnection struct { + NodeID string `gorm:"primaryKey;size:36" json:"node_id"` + // Empty when nobody holds this tunnel. A departed row keeps the empty + // string rather than NULL so there is one spelling of "held by nobody"; + // connectionIsHeld is the only place that spelling is written down. + OwnerInstanceID string `gorm:"size:36;index;not null" json:"owner_instance_id"` + Epoch int64 `gorm:"not null" json:"epoch"` + // No column DEFAULT: now() is PostgreSQL syntax and would reach the DDL, + // which breaks AutoMigrate on the SQLite single-binary path. Claim writes + // the database clock as an expression instead, the way Register does. + ConnectedAt time.Time `gorm:"not null" json:"connected_at"` + // DisconnectedAt is when the tunnel LEFT, and it exists because "gone for + // thirty seconds" and "never here" have to be different answers. + // + // A released row used to be deleted, so a worker re-homing between replicas + // was indistinguishable from one that had never dialled, and any grace + // period built on top would have had nothing to measure from. It is NOT a + // liveness clock for the owner: whether the OWNING replica is alive is + // still Instance.LastSeen and nothing else. It ticks once, on departure. + // + // Every writer in this package leaves it null while the row is held: Claim + // clears it in the statement that writes the owner, and only a departure + // sets it. That is a property of these writers, not of the table, and a + // mixed-version deployment breaks it: a replica running a binary from + // before this column existed claims without clearing the stamp, so a held + // row can carry the departure a newer replica recorded. + // + // So held-ness is the question, and the stamp only refines it. Ask + // connectionIsHeld first and read this second; a reader that reads the + // stamp alone reports a connected worker as gone. + DisconnectedAt *time.Time `gorm:"index" json:"disconnected_at,omitempty"` +} + +// connectionIsHeld is the one predicate that separates a row recording a tunnel +// somebody holds from a row recording a departure. Everything that asks that +// question asks it here, or asks for its negation: Owner refuses a +// row it rejects, ReapStale clears only rows it accepts, and +// PurgeDepartedBefore deletes only rows it rejects. Two spellings of one fact +// drift, and the drift here would show up as a departed worker resolving to a +// replica that holds nothing. +// +// The writes ask it too, and for the same reason Owner does rather than leaning +// on its join: Release and Deregister already match an owner id, but that only +// excludes a departed row while no owner id is ever empty, which is an accident +// of who registers. An empty one would match every departure in the table and +// reset its age. +// +// The column is table-qualified because Owner reads it across a join, where an +// unqualified name is ambiguous; qualifying it costs the single-table readers +// nothing. +const connectionIsHeld = `node_connections.owner_instance_id <> ''` + +// Migrate creates every table and sequence this package owns. It is the one +// call a caller has to remember: gorm's AutoMigrate models tables and columns +// but has no notion of a sequence, and the connection fence draws its epochs +// from one, so a caller that knew only about AutoMigrate would leave a schema +// that looks complete and cannot claim. Safe to call repeatedly. +// +// It does not take the migration advisory lock itself. The caller holds it +// across every table in the deployment, and taking a second one here would +// either nest inside that one or, worse, be the reason someone stops holding +// the outer one. +func Migrate(ctx context.Context, db *gorm.DB) error { + if err := db.WithContext(ctx).AutoMigrate(&Instance{}, &NodeConnection{}); err != nil { + return fmt.Errorf("migrating cluster tables: %w", err) + } + return ensureEpochSequence(ctx, db) +} + +// ensureEpochSequence creates the sequence Claim draws epochs from. It lives +// here, beside the model that needs it, because gorm's AutoMigrate models +// tables and columns but has no notion of a sequence; the caller that owns the +// migration advisory lock calls Migrate so that concurrently starting replicas +// do not race on the DDL. It is safe to call repeatedly. +// +// The sequence is not attached as a column DEFAULT on purpose: AutoMigrate +// compares the struct's declared default against the one PostgreSQL reports +// (`nextval('...'::regclass)`), and a mismatch there makes every startup ALTER +// the column. Naming the sequence in the statement keeps the schema stable. +func ensureEpochSequence(ctx context.Context, db *gorm.DB) error { + // CREATE SEQUENCE is PostgreSQL-only, and the same migration path runs + // against SQLite in single-binary mode. Nothing there can claim a + // connection (Claim refuses the dialect outright), so there is nothing to + // create. + if !isPostgres(db) { + return nil + } + if err := db.WithContext(ctx).Exec(`CREATE SEQUENCE IF NOT EXISTS ` + epochSequence + ` AS bigint`).Error; err != nil { + return fmt.Errorf("creating connection epoch sequence: %w", err) + } + return nil +} + +// Claim records ownerID as the owner of nodeID's tunnel and returns the epoch +// of the claim. +// +// The epoch is UNIQUE and never reissued: no other claim, for this node or any +// other, is ever handed the same value. It is NOT ordered, and callers must not +// treat it as a version number. The sequence value on the insert path is drawn +// while the tuple is built, before the row lock, so a claim that inserts after +// a Release can be handed a number lower than one already issued elsewhere. +// Compare epochs for equality only; never compare them for order. +// +// Uniqueness is all the fence needs: Release matches owner and epoch exactly, +// so a stale claim's token cannot match a live claim's row whichever way the +// two numbers happen to compare. +// +// It is one statement on purpose. A read-then-write would let two replicas read +// the same epoch and hand out the same fence token, which is exactly the case +// the fence exists to rule out; PostgreSQL serializes concurrent +// INSERT ... ON CONFLICT DO UPDATE on the conflicting row, so the losing writers +// block until the winner commits and only then draw their own epoch, in the +// order they took the row lock. +func (r *Registry) Claim(ctx context.Context, nodeID, ownerID string) (int64, error) { + // Refused rather than attempted on a dialect with no sequence. The + // statement would fail anyway, but with a driver-level "no such function: + // nextval" that reads like a missing migration; and a fence that cannot + // issue a token must not look like one that did. + if !isPostgres(r.db) { + return 0, fmt.Errorf("claiming connection for node %q as %q: connection ownership requires PostgreSQL, this deployment runs on %q", nodeID, ownerID, r.db.Dialector.Name()) + } + // connected_at is stamped by the database, never by this process, for the + // same reason instance liveness is: it is compared across replicas, so it + // has to be measured on the one clock every replica shares. + nextEpoch := gorm.Expr("nextval('" + epochSequence + "')") + values := map[string]any{ + "node_id": nodeID, + "owner_instance_id": ownerID, + "epoch": nextEpoch, + "connected_at": gorm.Expr("now()"), + } + if err := r.db.WithContext(ctx).Model(&NodeConnection{}).Clauses( + clause.OnConflict{ + Columns: []clause.Column{{Name: "node_id"}}, + DoUpdates: clause.Assignments(map[string]any{ + "owner_instance_id": ownerID, + "epoch": nextEpoch, + "connected_at": gorm.Expr("now()"), + // In the SAME statement as the owner, so a reconnect is never + // observed half-applied: a row that names a holder while still + // carrying a departure would answer "held" and "gone" at once. + "disconnected_at": nil, + }), + }, + clause.Returning{Columns: []clause.Column{{Name: "epoch"}}}, + ).Create(values).Error; err != nil { + return 0, fmt.Errorf("claiming connection for node %q as %q: %w", nodeID, ownerID, err) + } + // gorm scans RETURNING back over the map it was handed. If that ever stops + // happening the entry is still the expression we passed in, and returning a + // bogus epoch would hand out a fence token the database never issued. + epoch, ok := values["epoch"].(int64) + if !ok { + return 0, fmt.Errorf("claiming connection for node %q as %q: epoch not returned by the database (got %T)", nodeID, ownerID, values["epoch"]) + } + return epoch, nil +} + +// Owner returns the replica that holds nodeID's tunnel AND is still live, with +// the epoch of that claim, or ErrNoConnection when there is no such replica. +// +// This is the read anything that ACTS on the answer must use. A connection row +// outlives its owner: a replica that dies stops heartbeating but its rows stay +// until a peer's sweep removes them, which is up to InstanceLiveness plus one +// InstanceHeartbeat later. For that whole window the row names a process that +// is gone, and a relay built on it would dial a corpse and report the worker as +// unreachable rather than as absent. +// +// A missing row, a departed row and a dead owner are one answer on purpose. All +// three mean "no live replica holds this worker's tunnel", which is what a +// caller decides on; they differ only in which sweep has already run, and that +// is the sweeper's business rather than the caller's. +// +// That answer is emphatically not "this worker is gone". How long ago the +// departure was recorded is what separates a worker re-homing between replicas +// from one that has left, and this read does not report it: a caller that has +// to tell those apart reads the departure, not this. +// +// One statement, joined, not a row read followed by an instance lookup: between +// two statements the owner can die, and the caller would act on an owner the +// second read would have rejected. The join makes the two facts one snapshot. +// +// The window is InstanceLiveness rather than a parameter, which is the window +// the membership loop sweeps with. A caller free to pick its own could keep +// relaying to a replica the sweeper has already declared dead, or give up on +// one the sweeper is still keeping. +func (r *Registry) Owner(ctx context.Context, nodeID string) (string, int64, error) { + // Refused rather than attempted, for the reason Claim refuses: now() and + // make_interval are PostgreSQL, so on the single-binary SQLite path this + // would fail with "no such function: now", which reads as a missing + // migration. It is deliberately not ErrNoConnection. A deployment with no + // cluster has no answer to give about who owns a tunnel, and reporting + // absence would let a caller conclude the worker is not connected. + if !isPostgres(r.db) { + return "", 0, fmt.Errorf("looking up live owner of node %q: connection ownership requires PostgreSQL, this deployment runs on %q", nodeID, r.db.Dialector.Name()) + } + var conn NodeConnection + err := r.db.WithContext(ctx). + Model(&NodeConnection{}). + // Not load-bearing: gorm already expands this model's own columns, + // table-qualified, when a join is present and nothing was selected + // (callbacks.BuildQuerySQL). Written out so the projection is a + // property of this query rather than of that behaviour, since the join + // is here to filter and the row scanned back must stay this table's. + Select("node_connections.*"). + Joins("JOIN instances ON instances.id = node_connections.owner_instance_id AND "+instanceIsLive, InstanceLiveness.Seconds()). + // The departure filter is this query's own, not the join's. An empty + // owner matches no instance today only because no replica registers + // under an empty id, which is an accident of who registers rather than + // a property of ownership. + Where("node_connections.node_id = ? AND "+connectionIsHeld, nodeID). + Take(&conn).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return "", 0, fmt.Errorf("looking up live owner of node %q: %w", nodeID, ErrNoConnection) + } + if err != nil { + return "", 0, fmt.Errorf("looking up live owner of node %q: %w", nodeID, err) + } + return conn.OwnerInstanceID, conn.Epoch, nil +} + +// departure is the single assignment that turns a held row into a departed one. +// Release writes it for one claim and the membership sweep writes it for every +// claim a dead replica held, and the two must not disagree about what a +// departure looks like: a row cleared without a stamp, or stamped without being +// cleared, is a state every reader here has no answer for. +// +// The timestamp is the database's, never this process's, for the same reason +// instance liveness is: it is compared across replicas, so it has to be +// measured on the one clock they share. +func departure() map[string]any { + return map[string]any{ + "owner_instance_id": "", + "disconnected_at": gorm.Expr("now()"), + } +} + +// Release drops the claim identified by ownerID and epoch, and records the +// departure: the row survives with no owner and a disconnected_at stamp. +// +// An UPDATE and not a DELETE, because the row is the only place a departure can +// be recorded. Deleting it made a worker re-homing between replicas look like +// one that had never connected, so nothing above could tell a two-second blip +// from a worker that is gone. +// +// Both ownerID and epoch are in the WHERE so a replica that has only just +// noticed its dead socket cannot touch the claim a later reconnect established +// elsewhere: it must not clear a live claim, and it must not stamp a departure +// onto one. A claim that is no longer the live one is reported as +// ErrNoConnection rather than silently ignored, because the caller learning it +// has been fenced out is the point. +func (r *Registry) Release(ctx context.Context, nodeID, ownerID string, epoch int64) error { + // Refused rather than attempted, for the reason Claim refuses: the + // departure is stamped with now(), which on the single-binary SQLite path + // fails as a missing function and reads as a missing migration. It is + // deliberately not ErrNoConnection: a deployment with no cluster holds no + // claims, and reporting a fenced-out release would tell the caller it lost + // one. + if !isPostgres(r.db) { + return fmt.Errorf("releasing connection for node %q held by %q: connection ownership requires PostgreSQL, this deployment runs on %q", nodeID, ownerID, r.db.Dialector.Name()) + } + // gorm reports no error when a Where matches nothing, so the miss has to be + // read off RowsAffected. + res := r.db.WithContext(ctx). + Model(&NodeConnection{}). + // The held-ness filter is not implied by the owner match: a departed row + // keeps its epoch, so a release naming an empty owner would match it and + // stamp a fresh departure over the old one, making a worker that left + // long ago look like one that has only just gone. + Where("node_id = ? AND owner_instance_id = ? AND epoch = ? AND "+connectionIsHeld, nodeID, ownerID, epoch). + Updates(departure()) + if res.Error != nil { + return fmt.Errorf("releasing connection for node %q held by %q at epoch %d: %w", nodeID, ownerID, epoch, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("releasing connection for node %q held by %q at epoch %d: %w", nodeID, ownerID, epoch, ErrNoConnection) + } + return nil +} + +// PurgeDepartedBefore deletes the connection rows whose departure is older than +// olderThan, and returns how many went. +// +// Departures are kept so that absence can be decided, not kept forever. The +// retention has to outlast every window measured from a departure by enough +// that a purge can never turn a worker inside its reconnect grace into a worker +// that was never here; the caller passes a multiple of that grace, never the +// grace itself. +// +// A held row is never touched, whatever timestamp it carries: deleting one +// strands a worker that is connected at that moment. The age is measured by the +// database, like every other window here, so no replica's clock decides how old +// another replica's departure is. +func (r *Registry) PurgeDepartedBefore(ctx context.Context, olderThan time.Duration) (int64, error) { + // Refused rather than attempted, for the reason Claim refuses: now() and + // make_interval are PostgreSQL, and on the single-binary SQLite path this + // would fail with "no such function: now", which reads as a missing + // migration. + if !isPostgres(r.db) { + return 0, fmt.Errorf("purging departed connections: connection ownership requires PostgreSQL, this deployment runs on %q", r.db.Dialector.Name()) + } + res := r.db.WithContext(ctx). + Where("NOT ("+connectionIsHeld+") AND node_connections.disconnected_at < now() - make_interval(secs => ?)", + olderThan.Seconds()). + Delete(&NodeConnection{}) + if res.Error != nil { + return 0, fmt.Errorf("purging departed connections: %w", res.Error) + } + return res.RowsAffected, nil +} + +// connectedAmongQuery answers, for a set of nodes at once, which of their +// tunnels a LIVE replica holds and which of those are held by one named +// replica. +// +// ONE statement rather than a read per id, for the reason Owner is one +// statement: between two reads an owning replica can die, and the answer would +// then be assembled from two different snapshots of the cluster. A selection +// built that way can name a node whose owner the second read would have +// rejected, which is the one thing this read exists to prevent. +// +// The join is what makes a DEAD owner's row invisible here. Without it a row +// outlives its owner by up to a liveness window plus a heartbeat, and every +// caller that acted on this answer would be sent at a process that is gone. +// +// Held-ness is asked in SQL and the departure is not asked at all, and that is +// the whole difference between this read and Presence. This one answers a +// ROUTING question ("can a request get there right now"), so a row nobody holds +// is simply not a candidate, whatever its age. Deciding whether a worker has +// GONE is Presence's job and needs the grace; nothing here may be read as +// absence, and nothing here reports any. +// +// The liveness window is computed by the DATABASE, like every other window in +// this package: it is compared across replicas, so a Go-side cutoff would make +// the effective window depend on each replica's clock skew. The test container +// shares this host's clock, so no behavioural spec can see the difference and +// the statement shape is pinned instead. +// +// The predicates are the package's own rather than copies, so two spellings of +// "live" or of "held" cannot drift into a node that one query calls connected +// and another calls gone. +// +// The bind order is the order the placeholders appear in the text: the owner, +// then the liveness window, then the node ids. +const connectedAmongQuery = ` +SELECT node_connections.node_id AS node_id, + (node_connections.owner_instance_id = ?) AS by_owner +FROM node_connections +JOIN instances + ON instances.id = node_connections.owner_instance_id + AND ` + instanceIsLive + ` +WHERE ` + connectionIsHeld + ` + AND node_connections.node_id IN ?` + +// ConnectedAmong returns the subset of nodeIDs whose tunnel a LIVE replica +// currently holds, and separately those held by owner. +// +// It takes ids rather than asking which nodes are agents, because this package +// is a leaf and node type lives in core/services/nodes. +// +// heldByOwner is a SUBSET of held, not an alternative to it. A caller that +// prefers the ids it owns falls back to the rest, and expressing that as two +// disjoint lists would make the fallback a union the caller had to build. +// +// The answer is a routing fact and NEVER an absence one. An id missing from +// both lists means no live replica holds its tunnel at this instant, which a +// worker reconnecting between replicas produces routinely; nothing may read it +// as the worker having gone away. See Presence for the read that can answer +// that question, and only that one. +func (r *Registry) ConnectedAmong(ctx context.Context, nodeIDs []string, owner string) (held []string, heldByOwner []string, err error) { + if len(nodeIDs) == 0 { + // Answered without a statement, because `IN ()` is a syntax error in + // PostgreSQL and gorm's expansion of an empty slice is a filter nobody + // wrote. An empty candidate set has exactly one honest answer. + return nil, nil, nil + } + // Refused rather than attempted, for the reason Owner refuses: now() and + // make_interval are PostgreSQL, so on the single-binary SQLite path this + // would fail as a missing function and read as a missing migration. It is + // deliberately not an empty answer: a deployment with no cluster has + // nothing to say about who holds a tunnel, and an empty list would be read + // as "nobody does". + if !isPostgres(r.db) { + return nil, nil, fmt.Errorf("reading which of %d nodes are connected: connection ownership requires PostgreSQL, this deployment runs on %q", len(nodeIDs), r.db.Dialector.Name()) + } + var rows []struct { + NodeID string + ByOwner bool + } + if err := r.db.WithContext(ctx).Raw(connectedAmongQuery, + owner, InstanceLiveness.Seconds(), nodeIDs, + ).Scan(&rows).Error; err != nil { + return nil, nil, fmt.Errorf("reading which of %d nodes are connected: %w", len(nodeIDs), err) + } + for _, row := range rows { + held = append(held, row.NodeID) + if row.ByOwner { + heldByOwner = append(heldByOwner, row.NodeID) + } + } + return held, heldByOwner, nil +} diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go new file mode 100644 index 000000000000..00a8e40f6dd9 --- /dev/null +++ b/core/services/cluster/ownership_test.go @@ -0,0 +1,650 @@ +package cluster_test + +import ( + "context" + "fmt" + "path/filepath" + "regexp" + "strings" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" +) + +// sqlRecorder captures the statements gorm actually sends, so a spec can assert +// on the SQL rather than on gorm's intent. gorm silently drops clauses it +// cannot apply to a given destination, and such a drop turns an atomic upsert +// into something that still passes every sequential expectation. +type sqlRecorder struct { + gormlogger.Interface + mu sync.Mutex + statements []string + errs []error +} + +func newSQLRecorder() *sqlRecorder { + return &sqlRecorder{Interface: gormlogger.Default.LogMode(gormlogger.Silent)} +} + +func (r *sqlRecorder) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) { + sql, rows := fc() + r.mu.Lock() + r.statements = append(r.statements, sql) + if err != nil { + r.errs = append(r.errs, err) + } + r.mu.Unlock() + // Delegate so a failing statement is still reported the way gorm would + // report it. An instrument used to prove what the SQL does must not be the + // one thing that hides a statement erroring. + r.Interface.Trace(ctx, begin, func() (string, int64) { return sql, rows }, err) +} + +// only returns the single recorded statement, failing the spec if the call +// under test issued anything other than exactly one. +func (r *sqlRecorder) only() string { + r.mu.Lock() + defer r.mu.Unlock() + ExpectWithOffset(1, r.errs).To(BeEmpty(), "the recorded statement failed") + ExpectWithOffset(1, r.statements).To(HaveLen(1), "expected exactly one statement, got: %v", r.statements) + return r.statements[0] +} + +// statementCount reports how many statements were recorded, for the specs whose +// claim is that a call issued NONE. +func (r *sqlRecorder) statementCount() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.statements) +} + +// writeTarget matches the table a statement writes to, anchored at the verb so +// the UPDATE inside an upsert's ON CONFLICT clause cannot be mistaken for one. +var writeTarget = regexp.MustCompile(`^\s*(?i:delete\s+from|update)\s+"?([a-z_]+)"?`) + +// writeOrder returns the tables the recorded statements wrote to, in the order +// they were issued. It is how a spec pins a lock order: the order two paths +// take the same tables in is a property of the SQL, and asserting it on an +// outcome instead would mean racing two transactions into a real deadlock. +// +// Deletes and updates count alike. What deadlocks two transactions is the order +// they take row locks in, and an update takes the same lock a delete does, so a +// path that stopped deleting a table and started updating it would still have +// to keep the order. +func (r *sqlRecorder) writeOrder() []string { + r.mu.Lock() + defer r.mu.Unlock() + ExpectWithOffset(1, r.errs).To(BeEmpty(), "a recorded statement failed") + var tables []string + for _, stmt := range r.statements { + if m := writeTarget.FindStringSubmatch(stmt); m != nil { + tables = append(tables, m[1]) + } + } + return tables +} + +var _ = Describe("Connection ownership", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("hands every claim an epoch no other claim was given", func() { + // Uniqueness, not order. Claim's contract is that no two claims ever + // share an epoch; the insert path draws its sequence value before the + // row lock, so a claim that follows a Release can be handed a lower + // number than one already issued. Asserting e2 > e1 here would pin an + // ordering the fence does not need and does not promise. + e1, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + e2, err := reg.Claim(ctx, "w1", "inst-b") + Expect(err).ToNot(HaveOccurred()) + Expect(e2).ToNot(Equal(e1)) + }) + + It("reports the latest owner", func() { + _, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + e2, err := reg.Claim(ctx, "w1", "inst-b") + Expect(err).ToNot(HaveOccurred()) + + owner, epoch, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-b")) + Expect(epoch).To(Equal(e2), "the stored epoch must be the one the winning claim was handed") + }) + + It("distinguishes an unknown connection", func() { + _, _, err := ownerRow(ctx, db, "ghost") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("refuses a release from a stale owner", func() { + e1, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + _, err = reg.Claim(ctx, "w1", "inst-b") + Expect(err).ToNot(HaveOccurred()) + + // inst-a tries to clean up after losing the claim. + Expect(reg.Release(ctx, "w1", "inst-a", e1)).ToNot(Succeed()) + + owner, _, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-b"), "a stale owner must not be able to delete a live claim") + }) + + It("refuses a release that names the live owner but a stale epoch", func() { + // The same replica can reconnect a worker to itself; only the epoch + // separates the dead link from the live one. + e1, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + _, err = reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + Expect(reg.Release(ctx, "w1", "inst-a", e1)).ToNot(Succeed()) + + owner, _, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-a")) + }) + + It("never hands a node the same epoch twice, so a delayed cleanup cannot delete a live claim", func() { + // The scenario the fence exists for, with a release in the middle of it: + // inst-a claims and its link then dies silently. + eA1, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + // The worker reconnects to inst-b, which later releases cleanly. + eB, err := reg.Claim(ctx, "w1", "inst-b") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "inst-b", eB)).To(Succeed()) + // The worker comes back to inst-a, which is the same process throughout, + // so the owner id alone cannot separate this claim from the dead one. + eA2, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + // inst-a finally notices the first link is dead and cleans up after it. + // The harm is asserted before the cause, so a regression fails on the + // live claim disappearing rather than on the epoch arithmetic. + Expect(reg.Release(ctx, "w1", "inst-a", eA1)).ToNot(Succeed()) + + owner, epoch, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred(), "the delayed cleanup deleted the live claim") + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(eA2)) + Expect(eA2).ToNot(Equal(eA1), "an epoch handed out before a release must never be handed out again") + }) + + It("lets the current owner release its own claim", func() { + e, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "inst-a", e)).To(Succeed()) + + _, _, err = ownerRow(ctx, db, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + // Owner is the resolving read: it answers "who holds this tunnel and can be + // relayed to", where the raw row answers "what does the table say". The gap + // between the two is a whole liveness window wide, because a replica that + // dies leaves its connection rows behind until a peer's sweep removes them. + Describe("resolving the owner that can actually be relayed to", func() { + // Just past the window the membership loop sweeps with, not an arbitrary + // large age: a row aged ten minutes is rejected by any window between + // zero and ten minutes, so it would pin "filtered by SOME window" while + // letting Owner and the sweeper drift apart. The two seconds keep the + // spec off the exact boundary without loosening what it holds. + agedOut := cluster.InstanceLiveness + 2*time.Second + // Old enough that a narrowed window would reject it, still inside the + // one Owner must use. It is the other half of the same pin: agedOut + // fails a widened window, this fails a narrowed one. + agedButLive := cluster.InstanceLiveness / 2 + + // age rewrites an instance's heartbeat into the past. Sleeping for a + // liveness window is forbidden in a spec, and would be measuring the + // clock rather than the query. + age := func(id string, by time.Duration) { + ExpectWithOffset(1, db.Model(&cluster.Instance{}).Where("id = ?", id). + Update("last_seen", time.Now().Add(-by)).Error).To(Succeed()) + } + + It("names an owner whose replica is live", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + claimed, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + owner, epoch, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(claimed), "the resolved epoch must be the fence token the claim was handed") + }) + + It("still names an owner whose heartbeat is old but inside the window", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + claimed, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age("inst-a", agedButLive) + + owner, epoch, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred(), "a replica within the sweeper's window is alive, and its workers are still reachable through it") + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(claimed)) + }) + + It("refuses to name an owner that has no instance row at all", func() { + // What a completed sweep leaves for the moment between deleting the + // instance row and deleting the connections it orphaned, and what a + // re-registering replica's own connection rows look like meanwhile. + _, err := reg.Claim(ctx, "w1", "inst-gone") + Expect(err).ToNot(HaveOccurred()) + + _, _, err = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("refuses to name an owner whose heartbeat has aged past the liveness window", func() { + // The window this task exists to close: the replica is dead, no peer + // has swept it yet, and the row still names it. + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age("inst-a", agedOut) + + _, _, err = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("names an owner again once its heartbeat comes back", func() { + // Liveness is a window, not a latch: a replica that stalls and + // recovers still owns the sockets it never dropped, so resolution + // has to follow last_seen rather than remember a verdict. + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + claimed, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age("inst-a", agedOut) + Expect(reg.Heartbeat(ctx, "inst-a")).To(Succeed()) + + owner, epoch, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(claimed)) + }) + + It("still reports the dead owner in the table, which is why Owner reads liveness", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + claimed, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age("inst-a", agedOut) + + owner, epoch, err := ownerRow(ctx, db, "w1") + Expect(err).ToNot(HaveOccurred(), "the table holds the row and nothing else; hiding the dead owner here would leave the sweeper with no way to see what it has to clean up") + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(claimed)) + + _, _, err = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), "Owner must not agree with the raw row here") + }) + + It("reports a node with no connection at all the same way", func() { + _, _, err := reg.Owner(ctx, "ghost") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("resolves in one joined statement measured on the database clock", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + _, _, err = recording.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + + sql := strings.ToLower(rec.only()) + // only() rules out the read-then-look-up shape: two statements + // leave a window in which the owner dies between them, which is the + // race the join closes. + Expect(sql).To(ContainSubstring("join")) + Expect(sql).To(ContainSubstring("instances")) + // Liveness is compared across replicas, so the cutoff has to be + // computed on the one clock they all share. A Go-side time.Now() + // would appear as a bound parameter and a plain comparison instead, + // and replica clock skew would then widen or narrow the window. + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).To(ContainSubstring("make_interval")) + Expect(sql).ToNot(MatchRegexp(`last_seen\s*>\s*'`), + "the liveness cutoff must not be a literal timestamp from this process's clock") + }) + }) + + It("claims in one statement that draws its epoch from the database sequence and stamps on the database clock", func() { + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + + _, err := recording.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + sql := strings.ToLower(rec.only()) + // A read-then-write would show up here as two statements; the length + // check in only() is what rules that out. The rest pins the parts a + // silently dropped clause would remove. + Expect(sql).To(ContainSubstring("on conflict")) + Expect(sql).To(MatchRegexp(`(?i)nextval\s*\(\s*'node_connection_epochs'\s*\)`), + "the epoch must be drawn by the database, not computed by this process") + Expect(sql).To(ContainSubstring("returning")) + Expect(sql).To(ContainSubstring(`"epoch"`)) + // Timestamps are compared across replicas, so they must be measured on + // the one clock every replica shares. A Go-side time.Now() would appear + // as a bound parameter instead. + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).ToNot(MatchRegexp(`connected_at"?\s*=\s*'`), + "connected_at must not be a literal timestamp from this process's clock") + }) + + It("gives every concurrent claimant a distinct epoch and leaves exactly one winner", func() { + const claimants = 8 + epochs := make(chan int64, claimants) + var wg sync.WaitGroup + for i := 0; i < claimants; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + defer GinkgoRecover() + e, err := reg.Claim(context.Background(), "w-race", fmt.Sprintf("inst-%d", n)) + Expect(err).ToNot(HaveOccurred()) + epochs <- e + }(i) + } + wg.Wait() + close(epochs) + + seen := map[int64]bool{} + for e := range epochs { + Expect(seen[e]).To(BeFalse(), "epoch %d handed out twice; the fence is not atomic", e) + seen[e] = true + } + Expect(seen).To(HaveLen(claimants)) + + // Exactly one row, holding one of the epochs that was handed out: a + // winner, not a value nobody was given. Which of the eight wins is not + // asserted, and neither is any ordering among them. Epochs are unique + // and unordered, and a spec that ranked them here would teach the + // opposite of what Claim documents, whatever the sequence happens to do + // on this path. + var rows []cluster.NodeConnection + Expect(db.Where("node_id = ?", "w-race").Find(&rows).Error).To(Succeed()) + Expect(rows).To(HaveLen(1)) + Expect(seen).To(HaveKey(rows[0].Epoch), "the stored epoch was never handed to any claimant") + }) + + // A tunnel that goes away has to leave a mark. Deleting the row made "this + // worker's link dropped a moment ago" and "this worker has never connected + // here" one observation, so nothing above could tell a reconnect in flight + // from a departure, and every grace period built on top would have had + // nothing to measure from. + Describe("recording a departure", func() { + It("keeps the row and stamps disconnected_at when a claim is released", func() { + epoch, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + Expect(reg.Release(ctx, "w1", "inst-a", epoch)).To(Succeed()) + + var row cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&row).Error).To(Succeed(), + "the released row was deleted, so a worker that just left is indistinguishable from one that never dialled") + Expect(row.OwnerInstanceID).To(BeEmpty()) + Expect(row.DisconnectedAt).ToNot(BeNil()) + }) + + It("records the departure with one update measured on the database clock", func() { + epoch, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + Expect(recording.Release(ctx, "w1", "inst-a", epoch)).To(Succeed()) + + sql := strings.ToLower(rec.only()) + Expect(sql).To(ContainSubstring("update")) + Expect(sql).ToNot(ContainSubstring("delete")) + // The departure is compared against a grace window by other + // replicas, so it has to be stamped on the one clock they share. A + // Go-side time.Now() would appear as a bound parameter instead, and + // clock skew would then widen or narrow every window built on it. + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).ToNot(MatchRegexp(`disconnected_at"?\s*=\s*'`), + "the departure must not be a literal timestamp from this process's clock") + }) + + It("reports a released row as no connection from both reads", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "inst-a", epoch)).To(Succeed()) + + // Both reads: the row surviving is for whoever decides how old the + // departure is, and until something does, a departed row is no + // connection to a dialer and no connection to an observer. + _, _, err = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + _, _, err = ownerRow(ctx, db, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("does not resolve a departed row through an instance whose id is empty", func() { + // Owner rejects a departed row itself rather than leaning on the + // instances join to miss it. The join only misses an empty owner + // for as long as no instance row carries an empty id, which is an + // accident of who registers rather than a property of ownership. + Expect(reg.Register(ctx, "", "v1", 0, "")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "inst-a", epoch)).To(Succeed()) + + _, _, err = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("clears the departure when the worker reconnects", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "inst-a", epoch)).To(Succeed()) + + again, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(again).ToNot(Equal(epoch), "uniqueness, never ordering") + + var row cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&row).Error).To(Succeed()) + Expect(row.DisconnectedAt).To(BeNil(), + "a row with an owner still carried a departure, so a reader has two answers to choose from") + owner, _, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-a")) + }) + + It("does not let a fenced-out replica stamp a departure onto the live claim", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + stale, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + live, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + Expect(reg.Release(ctx, "w1", "inst-a", stale)).To(MatchError(cluster.ErrNoConnection)) + + owner, epoch, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(live)) + var row cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&row).Error).To(Succeed()) + Expect(row.DisconnectedAt).To(BeNil(), + "a replica that only just noticed its dead socket marked a live tunnel as departed") + }) + + It("refuses a release that names no owner, so a recorded departure cannot be aged backwards", func() { + // A departed row keeps the epoch of the claim that left, and its + // owner is the empty string, so a release naming an empty owner + // matches it on both columns. Stamping a fresh departure there would + // make a worker that left long ago look like one that has only just + // gone, which is the difference every window above is measured from. + epoch, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "w1", "inst-a", epoch)).To(Succeed()) + var before cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&before).Error).To(Succeed()) + + Expect(reg.Release(ctx, "w1", "", epoch)).To(MatchError(cluster.ErrNoConnection)) + + var after cluster.NodeConnection + Expect(db.WithContext(ctx).Where("node_id = ?", "w1").First(&after).Error).To(Succeed()) + Expect(after.DisconnectedAt).ToNot(BeNil()) + Expect(*after.DisconnectedAt).To(Equal(*before.DisconnectedAt)) + }) + + It("purges a departure older than the retention and keeps a recent one", func() { + oldEpoch, err := reg.Claim(ctx, "old", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "old", "inst-a", oldEpoch)).To(Succeed()) + // Aged on the database clock, which is the clock the purge measures + // on. Sleeping out a retention window in a spec is forbidden, and + // would be measuring this process's clock rather than the query. + Expect(db.WithContext(ctx).Exec( + `UPDATE node_connections SET disconnected_at = now() - make_interval(secs => ?) WHERE node_id = ?`, + 3600, "old").Error).To(Succeed()) + + freshEpoch, err := reg.Claim(ctx, "fresh", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "fresh", "inst-a", freshEpoch)).To(Succeed()) + + purged, err := reg.PurgeDepartedBefore(ctx, 10*time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(purged).To(Equal(int64(1))) + + var remaining []cluster.NodeConnection + Expect(db.WithContext(ctx).Find(&remaining).Error).To(Succeed()) + Expect(remaining).To(HaveLen(1)) + Expect(remaining[0].NodeID).To(Equal("fresh"), + "the purge took a departure that is still inside every window built on it") + }) + + It("measures the retention on the database clock", func() { + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + _, err := recording.PurgeDepartedBefore(ctx, 10*time.Minute) + Expect(err).ToNot(HaveOccurred()) + + sql := strings.ToLower(rec.only()) + // Asserted on the statement because no outcome can separate the two + // here: every replica in a spec shares this host's clock, so a + // cutoff computed in Go agrees with the database's until two + // machines disagree, and then one replica purges departures its + // peers still consider recent. + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).To(ContainSubstring("make_interval")) + Expect(sql).ToNot(MatchRegexp(`disconnected_at"?\s*<\s*'`), + "the retention cutoff must not be a literal timestamp from this process's clock") + }) + + It("never purges a row that is still held, whatever timestamp it carries", func() { + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + claimed, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + // No writer produces a held row carrying a departure, so it is + // written here directly: the purge has to be pinned on the owner + // column rather than on the timestamp happening to be null, because + // deleting a held row strands a worker that is connected right now. + Expect(db.WithContext(ctx).Exec( + `UPDATE node_connections SET disconnected_at = now() - make_interval(secs => ?) WHERE node_id = ?`, + 3600, "w1").Error).To(Succeed()) + + purged, err := reg.PurgeDepartedBefore(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(purged).To(BeZero()) + + owner, epoch, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred(), "the purge deleted the row of a tunnel somebody holds") + Expect(owner).To(Equal("inst-a")) + Expect(epoch).To(Equal(claimed)) + }) + }) +}) + +var _ = Describe("Connection ownership on a non-PostgreSQL dialect", func() { + var ( + db *gorm.DB + ctx context.Context + ) + + BeforeEach(func() { + var err error + ctx = context.Background() + db, err = gorm.Open(sqlite.Open(filepath.Join(GinkgoT().TempDir(), "cluster.db")), &gorm.Config{}) + Expect(err).ToNot(HaveOccurred()) + }) + + It("migrates, because the single-binary path shares this schema", func() { + // A PostgreSQL-only column DEFAULT here breaks AutoMigrate for every + // SQLite caller of nodes.NewNodeRegistry, which is how this regressed. + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + }) + + It("refuses to resolve an owner, rather than failing as a missing function", func() { + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + + _, _, err := cluster.NewRegistry(db).Owner(ctx, "w1") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("requires PostgreSQL")) + Expect(err.Error()).ToNot(ContainSubstring("no such function"), + "a dialect that cannot answer must say so, not surface as a missing migration") + Expect(err).ToNot(MatchError(cluster.ErrNoConnection), + "a deployment with no cluster has no answer about ownership; reporting absence would let a caller conclude the worker is not connected") + }) + + It("refuses to claim, rather than pretending to fence", func() { + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + + _, err := cluster.NewRegistry(db).Claim(ctx, "w1", "inst-a") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("requires PostgreSQL")) + }) + + It("refuses to release, rather than clearing a claim outside the fence", func() { + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + + err := cluster.NewRegistry(db).Release(ctx, "w1", "inst-a", 1) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("requires PostgreSQL")) + Expect(err.Error()).ToNot(ContainSubstring("no such function"), + "a dialect that cannot record a departure must say so, not surface as a missing migration") + Expect(err).ToNot(MatchError(cluster.ErrNoConnection), + "a deployment with no cluster holds no claims; reporting a fenced-out release would let a caller conclude it lost one") + }) + + It("refuses to purge departures, rather than failing as a missing function", func() { + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + + _, err := cluster.NewRegistry(db).PurgeDepartedBefore(ctx, time.Minute) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("requires PostgreSQL")) + Expect(err.Error()).ToNot(ContainSubstring("no such function"), + "a dialect that cannot answer must say so, not surface as a missing migration") + }) +}) diff --git a/core/services/cluster/presence.go b/core/services/cluster/presence.go new file mode 100644 index 000000000000..b0b8fdf0c379 --- /dev/null +++ b/core/services/cluster/presence.go @@ -0,0 +1,171 @@ +package cluster + +import ( + "context" + "fmt" + "time" +) + +// Presence is what this package can say about a worker's tunnel. +// +// Four values and not a boolean, because the conditions underneath are four and +// only ONE of them is a verdict a caller may act on. Collapsing any pair is the +// defect that this type exists to prevent: absence is what makes the scheduler stop +// placing work, reap node rows and evict models, and one of those paths runs +// during inference, so a worker misreported as absent loses the models it is +// serving at that moment. +// +// This type does not cover the two conditions that belong to the dialing path: +// a replica that would not answer, and the worker's own reply (including its +// "that backend is not there"). The dial code reports those, and they are not +// repeated here. +type Presence uint8 + +const ( + // PresenceUnknown reports that no connection row exists for this node at + // all. This package cannot tell a worker that has never dialled from one + // whose departure aged out of retention, and it must not guess: the first + // is a worker still starting up and the second is one long gone. + // + // It is the zero value on purpose, so a Presence returned alongside an + // error is the value nobody may act on rather than one that reads as a + // verdict. + PresenceUnknown Presence = iota + // PresenceConnected reports that a LIVE replica holds this worker's tunnel. + PresenceConnected + // PresenceReconnecting reports that no live replica holds the tunnel and + // the grace has not run out: either the departure is recent, or the owning + // replica died and no departure has been stamped yet. A retry, never a + // verdict; nobody may act on it. + PresenceReconnecting + // PresenceGone reports that no live replica holds the tunnel and the + // departure is older than the grace. This is the ONLY value a caller may + // read as absence. + PresenceGone +) + +// String names the value for a log line. The default case names the number +// rather than falling through to a real value: a caller reading "gone" out of a +// log for a value that does not exist would be reading the one answer it is +// allowed to reap on. +func (p Presence) String() string { + switch p { + case PresenceUnknown: + return "unknown" + case PresenceConnected: + return "connected" + case PresenceReconnecting: + return "reconnecting" + case PresenceGone: + return "gone" + default: + return fmt.Sprintf("presence(%d)", uint8(p)) + } +} + +// presenceQuery answers all three questions about one node in one statement: +// does a row exist, does a LIVE replica hold it, and is its departure older +// than the grace. +// +// One statement rather than a read followed by an instance lookup, for the +// reason Owner is one statement: between two reads the owning replica can die, +// and the answer would then be assembled from two different snapshots of the +// cluster. A dialer built that way relays into a dead replica for a whole +// liveness window. +// +// HELD-NESS IS ASKED FIRST, and the departure only refines it. Every writer in +// this binary clears disconnected_at in the same statement that writes the +// owner, but that is a property of these writers rather than of the table: a +// replica running a binary from before the column existed re-claims WITHOUT +// clearing the stamp, so during a rolling upgrade a HELD row can carry a +// departure older than any grace. Reading the stamp first, or treating a +// non-null stamp as evidence of absence on its own, reports a worker that is +// connected right now as gone. That is why `departed` carries the negation of +// connectionIsHeld rather than standing on the timestamp alone. +// +// Both windows are computed by the DATABASE. They are compared across replicas, +// so a Go-side cutoff would make the effective window depend on each replica's +// clock skew, and replicas disagreeing about whether a worker is gone is the +// flapping that this type exists to remove. No behavioural spec can see the +// difference (the test container shares the host clock), which is why the +// statement shape is pinned instead. +// +// The predicates are the package's own, not copies. Two spellings of "live" or +// of "held" drift, and the drift here reads as a worker that one query calls +// connected and another calls gone. +// +// The bind order is the order the placeholders appear in the text: the grace, +// then the liveness window, then the node id. +const presenceQuery = ` +SELECT + (` + connectionIsHeld + ` AND instances.id IS NOT NULL) AS held, + (NOT (` + connectionIsHeld + `) + AND node_connections.disconnected_at IS NOT NULL + AND node_connections.disconnected_at < now() - make_interval(secs => ?)) AS departed +FROM node_connections +LEFT JOIN instances + ON instances.id = node_connections.owner_instance_id + AND ` + instanceIsLive + ` +WHERE node_connections.node_id = ?` + +// Presence reports what this deployment can say about nodeID's tunnel, with +// grace as the window a departure has to outlive before it becomes a verdict. +// +// The grace is a parameter rather than a constant because it is an operator's +// trade (the worker reconnect grace), where the liveness window +// inside the query is not: that one has to be the window the membership loop +// sweeps with, or a reader would keep trusting an owner the sweeper has already +// declared dead. +// +// A held row whose owner is no longer live is PresenceReconnecting and not +// PresenceGone, deliberately. Nothing has stamped a departure on it yet, so +// there is no age to compare and the grace clock has not started; the +// membership sweep is what starts it. A replica dying must not condemn every +// worker it held. +func (r *Registry) Presence(ctx context.Context, nodeID string, grace time.Duration) (Presence, error) { + // Refused rather than attempted, for the reason Owner refuses: now() and + // make_interval are PostgreSQL, so on the single-binary SQLite path this + // would fail with "no such function: now", which reads as a missing + // migration. The refusal carries PresenceUnknown, the one value nobody may + // act on: a deployment with no cluster has no answer about a worker's + // tunnel, and reporting PresenceGone would license a reap. + if !isPostgres(r.db) { + return PresenceUnknown, fmt.Errorf("reading presence of node %q: connection ownership requires PostgreSQL, this deployment runs on %q", nodeID, r.db.Dialector.Name()) + } + var row struct { + Held bool + Departed bool + } + // gorm's Scan leaves the destination zero-valued when nothing matched and + // reports no error, so "no row" and "a row whose booleans are both false" + // are told apart by RowsAffected and by nothing else. Scan does not produce + // gorm.ErrRecordNotFound, so matching on that would silently turn every + // missing row into PresenceReconnecting. + res := r.db.WithContext(ctx).Raw(presenceQuery, + grace.Seconds(), InstanceLiveness.Seconds(), nodeID, + ).Scan(&row) + if res.Error != nil { + return PresenceUnknown, fmt.Errorf("reading presence of node %q: %w", nodeID, res.Error) + } + if res.RowsAffected == 0 { + return PresenceUnknown, nil + } + switch { + case row.Held: + // First, and redundantly so: the query already excludes held rows from + // `departed` (on the RAW column, before the join), so either gate alone + // answers Connected for a held row with a live owner, and no + // behavioural spec can tell which one is doing the work. + // + // The redundancy is PARTIAL, not a second complete gate. With the SQL + // gate removed and the owner DEAD, `held` is false and `departed` is + // true, so this ordering yields Gone where Reconnecting is correct. + // That is why the SQL gate carries its own assertion in the + // statement-shape spec: removing it is caught there, not here. + return PresenceConnected, nil + case row.Departed: + return PresenceGone, nil + default: + return PresenceReconnecting, nil + } +} diff --git a/core/services/cluster/presence_test.go b/core/services/cluster/presence_test.go new file mode 100644 index 000000000000..ae84a925cbbc --- /dev/null +++ b/core/services/cluster/presence_test.go @@ -0,0 +1,304 @@ +package cluster_test + +import ( + "context" + "math" + "path/filepath" + "strings" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +// age pushes a departure or a heartbeat into the past ON THE DATABASE CLOCK. +// Writing a Go-side time.Now().Add(-d) would age the row against this process's +// clock and then compare it against the database's, which is precisely the skew +// every window in this package is written to be immune to; a helper that did +// that would make the specs agree with an implementation that has the bug. +func age(ctx context.Context, db *gorm.DB, table, column, keyColumn, key string, by time.Duration) { + GinkgoHelper() + Expect(db.WithContext(ctx).Exec( + `UPDATE `+table+` SET `+column+` = now() - make_interval(secs => ?) WHERE `+keyColumn+` = ?`, + by.Seconds(), key).Error).To(Succeed()) +} + +var _ = Describe("Presence", func() { + var ( + ctx context.Context + db *gorm.DB + reg *cluster.Registry + ) + + // The grace every spec measures against, so the ageing below can sit just + // over the edge of it rather than an order of magnitude past it: a window + // only ever aged ten times its own width is a window nothing pins, and + // widening it stays green. + const grace = 60 * time.Second + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) + }) + + It("reports unknown for a node with no connection row", func() { + // Not "gone". This package cannot tell a worker that has never dialled + // from one whose departure aged out of retention, and a caller that + // reaps on absence would reap a worker that is still starting up. + p, err := reg.Presence(ctx, "never-seen", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceUnknown)) + }) + + It("reports connected while a live replica holds the tunnel", func() { + _, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceConnected)) + }) + + It("reports connected for a held row that still carries an old departure stamp", func() { + // The rolling-upgrade case, and the one thing this read must not get + // wrong. Every writer in THIS binary clears disconnected_at in the same + // statement that writes the owner, but a replica running a binary from + // before the column existed re-claims WITHOUT clearing it, so a held row + // can carry a departure older than any grace. Held-ness is the question + // and the stamp only refines it; a read that consults the stamp first, + // or treats a non-null stamp as evidence of absence, reports a worker + // that is connected right now as gone, for the whole upgrade. + _, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age(ctx, db, "node_connections", "disconnected_at", "node_id", "node-1", 10*grace) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceConnected), + "a row whose connection is HELD is present whatever its departure stamp says") + }) + + It("reports reconnecting for a held row whose owner is dead and whose stamp is stale", func() { + // The full mixed-version rolling-upgrade state, which no other spec + // constructs: the row is HELD, the owner is NOT live, and the stamp is + // older than the grace. It is reachable exactly once, during an upgrade + // where a pre-column replica re-claimed without clearing the stamp and + // then died before the sweep reached it. + // + // Reconnecting is the only defensible answer. The worker is re-dialling + // the load balancer right now, and the stamp on this row records a + // departure from a DIFFERENT, earlier session, so its age says nothing + // about the current one; the grace clock starts when the sweep stamps a + // departure for THIS session, and it has not run yet. + // + // Both of the ways this task can be got wrong land on a different + // answer here, which is why the state is worth constructing rather than + // reasoning about. Reading the stamp first gives Gone. Reading + // held-ness without the liveness join gives Connected. + _, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age(ctx, db, "node_connections", "disconnected_at", "node_id", "node-1", 10*grace) + age(ctx, db, "instances", "last_seen", "id", "inst-a", cluster.InstanceLiveness+5*time.Second) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceReconnecting), + "a held row with a dead owner is a retry; its stale stamp belongs to an earlier session and must not become a verdict") + }) + + It("reports reconnecting immediately after the tunnel is released", func() { + epoch, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "node-1", "inst-a", epoch)).To(Succeed()) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceReconnecting), + "a departure inside the grace is a retry, not a verdict; nobody may act on it") + }) + + It("reports gone once the departure is older than the grace", func() { + epoch, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "node-1", "inst-a", epoch)).To(Succeed()) + // One second past the edge, deliberately. Ageing by ten graces would + // leave a widened window green, and a window nothing pins is the defect + // class this spec exists for. + age(ctx, db, "node_connections", "disconnected_at", "node_id", "node-1", grace+time.Second) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceGone)) + }) + + It("still reports reconnecting one second inside the grace", func() { + // The other edge. Without it, a read that treats any departure as gone + // passes the spec above and condemns every worker re-homing normally. + epoch, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Release(ctx, "node-1", "inst-a", epoch)).To(Succeed()) + age(ctx, db, "node_connections", "disconnected_at", "node_id", "node-1", grace-time.Second) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceReconnecting)) + }) + + It("reports reconnecting, not gone, when the OWNING REPLICA dies and the row is still held", func() { + // A replica dying must not condemn every worker it held. The row still + // names inst-a and inst-a is no longer live, so nothing holds the + // tunnel; but no departure has been stamped, so the grace clock has not + // started and there is no age to compare. That is a retry. + _, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age(ctx, db, "instances", "last_seen", "id", "inst-a", cluster.InstanceLiveness+5*time.Second) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceReconnecting)) + }) + + It("reports gone after the membership sweep has stamped the departure and the grace has passed", func() { + // The full path a dead replica's worker takes: the sweep records the + // departure, and only then does the grace start running. + _, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age(ctx, db, "instances", "last_seen", "id", "inst-a", cluster.InstanceLiveness+5*time.Second) + + _, connections, err := reg.ReapStale(ctx, "inst-sweeper", cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(connections).To(Equal(int64(1)), "the sweep must have stamped the departure this spec then ages") + age(ctx, db, "node_connections", "disconnected_at", "node_id", "node-1", grace+time.Second) + + p, err := reg.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + Expect(p).To(Equal(cluster.PresenceGone)) + }) + + It("answers in one joined statement measured on the database clock", func() { + _, err := reg.Claim(ctx, "node-1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + + rec := newSQLRecorder() + recording := cluster.NewRegistry(db.Session(&gorm.Session{Logger: rec})) + _, err = recording.Presence(ctx, "node-1", grace) + Expect(err).ToNot(HaveOccurred()) + + sql := strings.ToLower(rec.only()) + // only() rules out the read-then-look-up shape: between two statements + // the owning replica can die, and the answer would then be built from + // two different snapshots of the cluster. + Expect(sql).To(ContainSubstring("join")) + Expect(sql).To(ContainSubstring("instances")) + // The departure test is gated on held-ness IN THE SQL, and not only by + // the order of the switch that reads these two booleans. Either gate + // alone gives the right answer today, so no behavioural spec can tell + // which one is doing the work, and a single-layer regression would sit + // here unnoticed until the second layer moved too. Held-ness first is + // the invariant that the absence rules rest on, so both layers are pinned. + Expect(strings.Join(strings.Fields(sql), " ")).To(ContainSubstring( + "not (node_connections.owner_instance_id <> '') and node_connections.disconnected_at is not null")) + // Both windows are computed by the database. A Go-side comparison would + // appear here as a bound literal, and replica clock skew would then move + // the grace and the liveness window per replica. The test container + // shares this host's clock, so no behavioural spec can catch that; the + // statement shape is the only place it is visible. + Expect(strings.Count(sql, "make_interval")).To(Equal(2), + "both the liveness window and the reconnect grace must be computed by the database") + Expect(sql).To(ContainSubstring("now()")) + Expect(sql).ToNot(MatchRegexp(`disconnected_at\s*<\s*'`), + "the grace cutoff must not be a literal timestamp from this process's clock") + Expect(sql).ToNot(MatchRegexp(`last_seen\s*>\s*'`), + "the liveness cutoff must not be a literal timestamp from this process's clock") + }) + + Describe("naming the values", func() { + It("gives every value a name a log line can carry", func() { + Expect(cluster.PresenceUnknown.String()).To(Equal("unknown")) + Expect(cluster.PresenceConnected.String()).To(Equal("connected")) + Expect(cluster.PresenceReconnecting.String()).To(Equal("reconnecting")) + Expect(cluster.PresenceGone.String()).To(Equal("gone")) + }) + + It("does not name an unknown value after a real one", func() { + // A default branch that fell through to "gone" would put the one + // value callers may act on into a log line for a value that does + // not exist. + Expect(cluster.Presence(200).String()).To(ContainSubstring("200")) + }) + }) + + Describe("the retention that outlives the grace", func() { + // The retention is not pinned to any one grace value on purpose. It is + // the purge's window and the grace is a reader's window, and the only + // property that matters is that the first always outlasts the second: + // a purge that deletes a departure a reader is still measuring against + // turns a worker inside its reconnect grace back into a worker that was + // never here, and an operator raising the grace past a FIXED retention + // would leave a genuinely gone worker reading as unknown forever, so + // nothing ever reaps it. + DescribeTable("outlasts the grace a reader measures against", + func(grace time.Duration) { + Expect(cluster.DepartedRetentionFor(grace)).To(BeNumerically(">", grace)) + }, + Entry("a grace of one second", time.Second), + Entry("the default grace", 60*time.Second), + Entry("a grace at the retention floor", cluster.DepartedRetention), + Entry("a grace far past the floor", 10*cluster.DepartedRetention), + ) + + It("does not wrap back to the floor for a grace so large the multiply overflows", func() { + // time.Duration is int64 nanoseconds, so five times a grace above + // roughly 58 years wraps. An unguarded multiply comes back negative + // or small, falls through to the floor, and reinstates exactly the + // defect this function exists to remove, silently and only for the + // operator who set the largest window. + huge := time.Duration(math.MaxInt64) + Expect(cluster.DepartedRetentionFor(huge)).To(BeNumerically(">=", huge)) + // And one just past the threshold, where the wrap is easiest to + // miss because the input still looks like an ordinary duration. + overflowing := time.Duration(math.MaxInt64/5) + time.Second + Expect(cluster.DepartedRetentionFor(overflowing)).To(BeNumerically(">=", overflowing)) + }) + + It("never shortens below the floor for a very small grace", func() { + // A tiny grace must not shrink the retention to match: the row is also + // what a restarting deployment reads to tell a re-dialling worker from + // one it has never seen. + Expect(cluster.DepartedRetentionFor(time.Second)).To(Equal(cluster.DepartedRetention)) + }) + }) +}) + +var _ = Describe("Presence on a non-PostgreSQL dialect", func() { + var ( + db *gorm.DB + ctx context.Context + ) + + BeforeEach(func() { + var err error + ctx = context.Background() + db, err = gorm.Open(sqlite.Open(filepath.Join(GinkgoT().TempDir(), "cluster.db")), &gorm.Config{}) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + }) + + It("refuses to answer, rather than failing as a missing function", func() { + p, err := cluster.NewRegistry(db).Presence(ctx, "node-1", time.Minute) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("requires PostgreSQL")) + Expect(err.Error()).ToNot(ContainSubstring("no such function"), + "a dialect that cannot answer must say so, not surface as a missing migration") + Expect(p).To(Equal(cluster.PresenceUnknown), + "a refusal must carry the value nobody may act on, never one that reads as absence") + }) +}) From 3e04ea1805ccaac35f2ab6f91f1437c404f21cd0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:15:27 +0000 Subject: [PATCH 009/162] feat(distributed): record which carriers a worker node is attached to A worker that follows a change of carrier is attached to two carriers for a while. Add the columns that hold what the worker reports and the epoch it saw when it reported. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/registry.go | 9 +++++++- core/services/nodes/registry_test.go | 34 ++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index f69d943685f3..f1f71b00bbd0 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -96,7 +96,14 @@ type BackendNode struct { // registration. Empty for workers registered before this field existed. Version string `gorm:"column:version;size:64" json:"version,omitempty"` // Commit is the git commit hash the worker binary was built from. - Commit string `gorm:"column:commit;size:64" json:"commit,omitempty"` + Commit string `gorm:"column:commit;size:64" json:"commit,omitempty"` + // Attached lists the carriers that the worker reported it is connected to, + // separated by commas, and AttachedEpoch is the epoch of the cluster + // carrier row that the worker saw when it reported. During a change of + // carrier a worker is attached to both, and the frontends route by this. + // Both are empty for a worker that does not report them. + Attached string `gorm:"column:attached;size:64" json:"attached,omitempty"` + AttachedEpoch int64 `gorm:"column:attached_epoch;default:0" json:"attached_epoch,omitempty"` APIKeyID string `gorm:"size:36" json:"-"` // auto-provisioned API key ID (for cleanup) AuthUserID string `gorm:"size:36" json:"-"` // auto-provisioned user ID (for cleanup) LastHeartbeat time.Time `gorm:"column:last_heartbeat" json:"last_heartbeat"` diff --git a/core/services/nodes/registry_test.go b/core/services/nodes/registry_test.go index 8d14323ec403..90955227b370 100644 --- a/core/services/nodes/registry_test.go +++ b/core/services/nodes/registry_test.go @@ -1911,3 +1911,37 @@ var _ = Describe("ModelScheduling spread + seeding", func() { Expect(n.SpreadAll).To(BeTrue()) }) }) + +var _ = Describe("NodeRegistry carrier attachment", func() { + var ( + db *gorm.DB + registry *NodeRegistry + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + }) + + It("stores no attachment for a node that reports none", func() { + node := &BackendNode{Name: "plain", NodeType: NodeTypeBackend, Address: "10.0.0.3:50051"} + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + + var got BackendNode + Expect(db.First(&got, "id = ?", node.ID).Error).To(Succeed()) + Expect(got.Attached).To(BeEmpty()) + Expect(got.AttachedEpoch).To(BeZero()) + }) + + It("keeps the attachment a node reported", func() { + node := &BackendNode{Name: "dual", NodeType: NodeTypeBackend, Address: "10.0.0.4:50051", Attached: "nats,tunnel", AttachedEpoch: 3} + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + + var got BackendNode + Expect(db.First(&got, "id = ?", node.ID).Error).To(Succeed()) + Expect(got.Attached).To(Equal("nats,tunnel")) + Expect(got.AttachedEpoch).To(Equal(int64(3))) + }) +}) From b8dce7b6c0e234171b7947cbeebe512435cc56e8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:15:27 +0000 Subject: [PATCH 010/162] feat(distributed): register every frontend replica in the cluster registry Run the membership loop in distributed mode, on every carrier, and remove the row when the replica stops. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 34 ++++++++++++++++++++ core/application/distributed_carrier_test.go | 29 ++++++++++++++++- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/core/application/distributed.go b/core/application/distributed.go index bb9b5845e05d..bef0ca7c4877 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -13,6 +13,7 @@ import ( "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" + "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/agents" "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" @@ -23,6 +24,7 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/xlog" @@ -61,6 +63,10 @@ type DistributedServices struct { // backend-logs proxy, the same way the HTTP file stager does. WorkerHTTPDial nodes.WorkerNetDialerFor + // Membership keeps the row of this replica in the instances table and + // sweeps the rows of replicas that stopped answering. + Membership *cluster.Membership + // active names the carrier set the holders forward to. active *atomic.Pointer[carrier.Set] @@ -74,6 +80,8 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { + // First, so the peers see this replica leave before its services stop. + ds.Membership.Stop() if ds.Health != nil { ds.Health.Stop() } @@ -533,8 +541,14 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Create ModelRouterAdapter to wire into ModelLoader modelAdapter := nodes.NewModelRouterAdapter(router) + membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID) + if err != nil { + return nil, err + } + success = true return &DistributedServices{ + Membership: membership, Broadcaster: broadcaster, WorkQueue: workQueue, AgentControl: carrier.NewAgents(active), @@ -560,6 +574,26 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade }, nil } +// startMembership creates the cluster tables, then records this replica in the +// instances table and keeps its row fresh until the Membership is stopped. +// Every replica does this on every carrier, because the change of carrier needs +// to know which replicas are alive and ready. +func startMembership(ctx context.Context, db *gorm.DB, instanceID string) (*cluster.Membership, error) { + if ctx == nil { + ctx = context.Background() + } + if err := advisorylock.WithLockCtx(ctx, db, advisorylock.KeySchemaMigrate, func() error { + return cluster.Migrate(ctx, db) + }); err != nil { + return nil, fmt.Errorf("migrating cluster tables: %w", err) + } + membership := cluster.NewMembership(cluster.NewRegistry(db), instanceID, internal.PrintableVersion()) + if err := membership.Start(ctx); err != nil { + return nil, fmt.Errorf("registering this replica: %w", err) + } + return membership, nil +} + func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go index b5609fb5dbb2..ad187afed869 100644 --- a/core/application/distributed_carrier_test.go +++ b/core/application/distributed_carrier_test.go @@ -35,7 +35,7 @@ var _ = Describe("distributed startup and the cluster carrier", func() { Distributed: config.DistributedConfig{ Enabled: true, InstanceID: "replica-a", - NatsURL: "nats://127.0.0.1:1", + NatsURL: "nats://127.0.0.1:1", }, } }) @@ -123,6 +123,33 @@ var _ = Describe("distributed startup and the cluster carrier", func() { Eventually(got).Should(Receive(Equal([]byte(`"hello"`)))) }) + It("registers every replica in the instances table and removes it on shutdown", func() { + live := func() []string { + GinkgoHelper() + instances, err := cluster.NewRegistry(db).Live(context.Background(), cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + ids := make([]string, 0, len(instances)) + for _, in := range instances { + ids = append(ids, in.ID) + } + return ids + } + + first, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(first.Shutdown) + cfg2 := *cfg + cfg2.Distributed.InstanceID = "replica-b" + cfg2.DataPath = GinkgoT().TempDir() + second, err := initDistributed(&cfg2, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Shutdown) + Expect(live()).To(ConsistOf("replica-a", "replica-b")) + + first.Shutdown() + Expect(live()).To(ConsistOf("replica-b")) + }) + It("starts a second replica on the same row without changing it", func() { first, err := initDistributed(cfg, db, nil, nil) Expect(err).ToNot(HaveOccurred()) From ad2ff0f87b134bc8db252c6a60bc79083c3ffd97 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:17:37 +0000 Subject: [PATCH 011/162] refactor(messaging): share the subject matcher and bound broadcast payloads Move the subject matcher out of the test helper into messaging, so a carrier and the in-memory doubles use one rule. Add MaxBroadcastBytes, checked by the NATS client and by the fake bus, so a message that one carrier accepts is not refused by another. The conformance suite now takes two ends of a carrier and runs every broadcast root, the sizes around the notify cap and the bound. It also asserts what a carrier does with the control roots, and reads the drop counter when the carrier has one. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../carrier/broadcaster_conformance_test.go | 10 +- .../galleryop/distributed_sync_test.go | 23 +-- core/services/messaging/client.go | 3 + .../messaging/client_conformance_test.go | 13 +- core/services/messaging/export_test.go | 9 + core/services/messaging/limits.go | 28 ++++ core/services/messaging/limits_test.go | 41 +++++ core/services/messaging/match.go | 38 +++++ core/services/messaging/match_test.go | 54 ++++++ .../messaging/messagingtest/broadcaster.go | 158 ++++++++++++++++-- core/services/testutil/export_test.go | 4 - core/services/testutil/fakebus.go | 29 +--- .../testutil/fakebus_conformance_test.go | 6 +- core/services/testutil/subject_match_test.go | 21 --- 14 files changed, 345 insertions(+), 92 deletions(-) create mode 100644 core/services/messaging/limits.go create mode 100644 core/services/messaging/limits_test.go create mode 100644 core/services/messaging/match.go create mode 100644 core/services/messaging/match_test.go delete mode 100644 core/services/testutil/export_test.go delete mode 100644 core/services/testutil/subject_match_test.go diff --git a/core/services/carrier/broadcaster_conformance_test.go b/core/services/carrier/broadcaster_conformance_test.go index 3d8b0e60345f..00cb57e0c34f 100644 --- a/core/services/carrier/broadcaster_conformance_test.go +++ b/core/services/carrier/broadcaster_conformance_test.go @@ -5,7 +5,6 @@ import ( "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" - "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/messaging/messagingtest" . "github.com/onsi/ginkgo/v2" ) @@ -14,15 +13,16 @@ import ( // the carriers do: first as it is built, then after its carrier was swapped. var _ = Describe("Broadcaster holder conformance", func() { Describe("on the carrier it was built over", func() { - messagingtest.RunBroadcasterConformance(func() (messaging.Broadcaster, func()) { + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { var cur atomic.Pointer[carrier.Set] cur.Store(newFakeCarrier(cluster.CarrierNATS, 1).set) - return carrier.NewBroadcaster(&cur), func() {} + h := carrier.NewBroadcaster(&cur) + return messagingtest.Carrier{Bus: h, Peer: h, ServesControlRoots: true} }) }) Describe("after the carrier was swapped", func() { - messagingtest.RunBroadcasterConformance(func() (messaging.Broadcaster, func()) { + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { var cur atomic.Pointer[carrier.Set] first := newFakeCarrier(cluster.CarrierNATS, 1).set next := newFakeCarrier(cluster.CarrierTunnel, 2).set @@ -35,7 +35,7 @@ var _ = Describe("Broadcaster holder conformance", func() { if err := h.Release(first); err != nil { Fail(err.Error()) } - return h, func() {} + return messagingtest.Carrier{Bus: h, Peer: h, ServesControlRoots: true} }) }) }) diff --git a/core/services/galleryop/distributed_sync_test.go b/core/services/galleryop/distributed_sync_test.go index b1c9a31202bc..862ca1565d35 100644 --- a/core/services/galleryop/distributed_sync_test.go +++ b/core/services/galleryop/distributed_sync_test.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "errors" - "strings" "sync" "time" @@ -37,26 +36,6 @@ type fakeBusSub struct { func newFakeBus() *fakeBus { return &fakeBus{} } -func subjectMatches(filter, subject string) bool { - if filter == subject { - return true - } - fp := strings.Split(filter, ".") - sp := strings.Split(subject, ".") - if len(fp) != len(sp) { - return false - } - for i := range fp { - if fp[i] == "*" { - continue - } - if fp[i] != sp[i] { - return false - } - } - return true -} - func (b *fakeBus) Publish(subject string, data any) error { payload, err := json.Marshal(data) if err != nil { @@ -66,7 +45,7 @@ func (b *fakeBus) Publish(subject string, data any) error { subs := append([]fakeBusSub(nil), b.subs...) b.mu.Unlock() for _, s := range subs { - if subjectMatches(s.subject, subject) { + if messaging.SubjectMatches(s.subject, subject) { s.handler(payload) } } diff --git a/core/services/messaging/client.go b/core/services/messaging/client.go index b8ab08059c16..7bce43c17918 100644 --- a/core/services/messaging/client.go +++ b/core/services/messaging/client.go @@ -154,6 +154,9 @@ func (c *Client) Publish(subject string, data any) error { if err != nil { return fmt.Errorf("marshalling message for %s: %w", subject, err) } + if err := CheckBroadcastSize(subject, len(payload)); err != nil { + return err + } c.mu.RLock() defer c.mu.RUnlock() return c.conn.Publish(subject, payload) diff --git a/core/services/messaging/client_conformance_test.go b/core/services/messaging/client_conformance_test.go index 8ab456710cff..e08d7171de84 100644 --- a/core/services/messaging/client_conformance_test.go +++ b/core/services/messaging/client_conformance_test.go @@ -60,7 +60,7 @@ var _ = AfterSuite(func() { }) var _ = Describe("NATS client", func() { - messagingtest.RunBroadcasterConformance(func() (messaging.Broadcaster, func()) { + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { url, err := sharedNATS() if err != nil { // This is the only spec that runs the rules against a real carrier. @@ -76,6 +76,15 @@ var _ = Describe("NATS client", func() { if err != nil { Fail("connecting to the test NATS server: " + err.Error()) } - return c, c.Close + peer, err := messaging.New(url) + if err != nil { + c.Close() + Fail("connecting to the test NATS server: " + err.Error()) + } + return messagingtest.Carrier{ + Bus: c, Peer: peer, + Cleanup: func() { c.Close(); peer.Close() }, + ServesControlRoots: true, + } }) }) diff --git a/core/services/messaging/export_test.go b/core/services/messaging/export_test.go index 8c68067524db..10a551b3e294 100644 --- a/core/services/messaging/export_test.go +++ b/core/services/messaging/export_test.go @@ -3,3 +3,12 @@ package messaging // NATSRouteForTest exposes the route table so the external specs can pin the // queue group per kind, which no producer-side behaviour reveals. var NATSRouteForTest = natsRoute + +// BroadcastRootsForTest returns the broadcast roots of the subject rules. +func BroadcastRootsForTest() []string { + roots := make([]string, 0, len(broadcastRoots)) + for r := range broadcastRoots { + roots = append(roots, r) + } + return roots +} diff --git a/core/services/messaging/limits.go b/core/services/messaging/limits.go new file mode 100644 index 000000000000..7f52f1416ef0 --- /dev/null +++ b/core/services/messaging/limits.go @@ -0,0 +1,28 @@ +package messaging + +import ( + "errors" + "fmt" +) + +// MaxBroadcastBytes is the largest encoded payload that a carrier accepts for a +// broadcast. Every carrier enforces the same bound, so a message that one +// carrier carries is not refused by another after a change of carrier. +// +// The bound is far above the payloads that the deployment sends. NATS refuses a +// payload above the max_payload of its server, which is 1 MiB unless an operator +// raised it, and this bound is for the carriers that have no such limit of their +// own. +const MaxBroadcastBytes = 8 << 20 + +// ErrPayloadTooLarge means a broadcast is larger than MaxBroadcastBytes. +var ErrPayloadTooLarge = errors.New("messaging: broadcast payload is too large") + +// CheckBroadcastSize returns an error that wraps ErrPayloadTooLarge when an +// encoded payload of size bytes is above MaxBroadcastBytes. +func CheckBroadcastSize(subject string, size int) error { + if size > MaxBroadcastBytes { + return fmt.Errorf("%w: %d bytes on %q, the bound is %d", ErrPayloadTooLarge, size, subject, MaxBroadcastBytes) + } + return nil +} diff --git a/core/services/messaging/limits_test.go b/core/services/messaging/limits_test.go new file mode 100644 index 000000000000..c3a3e1ed0c49 --- /dev/null +++ b/core/services/messaging/limits_test.go @@ -0,0 +1,41 @@ +package messaging_test + +import ( + "errors" + "maps" + "slices" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" +) + +var _ = Describe("The conformance suite and the subject rules", func() { + It("has a subject for exactly the broadcast roots that the rules serve", func() { + want := messaging.BroadcastRootsForTest() + slices.Sort(want) + got := slices.Sorted(maps.Keys(messagingtest.BroadcastRootSubjects)) + Expect(got).To(Equal(want), "add a row to messagingtest.BroadcastRootSubjects for each new broadcast root") + }) + + It("uses subjects that the rules accept, under the root that names them", func() { + for root, subject := range messagingtest.BroadcastRootSubjects { + Expect(messaging.ValidateSubject(subject)).To(Succeed(), subject) + Expect(subject).To(HavePrefix(root+"."), subject) + } + }) +}) + +var _ = Describe("CheckBroadcastSize", func() { + It("accepts the bound itself", func() { + Expect(messaging.CheckBroadcastSize("jobs.a.result", messaging.MaxBroadcastBytes)).To(Succeed()) + }) + + It("refuses one byte more, with a class that a caller can match", func() { + err := messaging.CheckBroadcastSize("jobs.a.result", messaging.MaxBroadcastBytes+1) + Expect(errors.Is(err, messaging.ErrPayloadTooLarge)).To(BeTrue()) + Expect(err.Error()).To(ContainSubstring("jobs.a.result")) + }) +}) diff --git a/core/services/messaging/match.go b/core/services/messaging/match.go new file mode 100644 index 000000000000..480ccc44c0ef --- /dev/null +++ b/core/services/messaging/match.go @@ -0,0 +1,38 @@ +package messaging + +import "strings" + +// SubjectMatches reports whether a subscription filter matches a concrete +// subject. A `*` that stands alone as a token matches exactly one token, as in +// NATS. +// +// It is the one definition of subject matching in the tree. A carrier that +// matches subjects itself and the in-memory doubles that specs publish through +// both call it, so a filter cannot match on one and miss on the other. +// +// A `>` tail wildcard is not implemented, and a filter with one matches +// nothing. A caller that writes one then receives no message and not every +// message on the prefix. The check comes before the test for equal strings, +// because otherwise `a.>` would match the literal subject `a.>`. +func SubjectMatches(filter, subject string) bool { + if strings.Contains(filter, ">") { + return false + } + if filter == subject { + return true + } + fp := strings.Split(filter, ".") + sp := strings.Split(subject, ".") + if len(fp) != len(sp) { + return false + } + for i := range fp { + if fp[i] == "*" { + continue + } + if fp[i] != sp[i] { + return false + } + } + return true +} diff --git a/core/services/messaging/match_test.go b/core/services/messaging/match_test.go new file mode 100644 index 000000000000..1eb731bfc5a8 --- /dev/null +++ b/core/services/messaging/match_test.go @@ -0,0 +1,54 @@ +package messaging_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// These rows are the contract between the carriers and the in-memory doubles +// that specs publish through. When each had its own copy of the rule, a filter +// could match on one and not on the other, and the difference looked like a peer +// that receives an event on one replica and misses it on another. +var _ = DescribeTable("SubjectMatches", + func(filter, subject string, expected bool) { + Expect(messaging.SubjectMatches(filter, subject)).To(Equal(expected)) + }, + Entry("an exact subject matches itself", "jobs.new", "jobs.new", true), + Entry("a different literal does not match", "jobs.new", "jobs.done", false), + Entry("matching is case sensitive", "Jobs.new", "jobs.new", false), + + Entry("two wildcards each consume one token", + "agent.*.events.*", "agent.myagent.events.user1", true), + Entry("a wildcard does not match a missing token", + "agent.*.events.*", "agent.myagent.events", false), + Entry("a wildcard does not swallow two tokens", + "agent.*.events.*", "agent.a.b.events.u", false), + Entry("a wildcard does not match a longer subject", + "agent.*.events.*", "agent.myagent.events.user1.extra", false), + + Entry("a middle wildcard matches one token", + "jobs.*.progress", "jobs.abc.progress", true), + Entry("a middle wildcard does not match two tokens", + "jobs.*.progress", "jobs.abc.def.progress", false), + Entry("a middle wildcard does not match zero tokens", + "jobs.*.progress", "jobs.progress", false), + + Entry("a bare wildcard matches exactly one token", "*", "jobs", true), + Entry("a bare wildcard does not match two tokens", "*", "jobs.new", false), + + // A model id that looks like several tokens is one token after the + // sanitizer, so the staging filter can be a single '*'. + Entry("the staging wildcard matches a sanitized multi-part model id", + messaging.SubjectStagingProgressWildcard, messaging.SubjectStagingProgress("a.b.c"), true), + + // '>' is not implemented. A caller that writes one must get nothing, and + // not every message on the prefix. + Entry("a tail wildcard filter matches nothing", "a.>", "a.b", false), + Entry("a tail wildcard filter matches nothing even one token deep", "a.>", "a.b.c", false), + Entry("a bare tail wildcard matches nothing", ">", "anything", false), + // The refusal comes before the check for equal strings. Without that order + // the filter would match the literal subject "a.>". + Entry("a tail wildcard filter does not even match itself", "a.>", "a.>", false), +) diff --git a/core/services/messaging/messagingtest/broadcaster.go b/core/services/messaging/messagingtest/broadcaster.go index b0d7f37cb49b..df3b158ad3e4 100644 --- a/core/services/messaging/messagingtest/broadcaster.go +++ b/core/services/messaging/messagingtest/broadcaster.go @@ -16,8 +16,51 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" ) -// Factory returns a ready carrier and a cleanup for it. -type Factory func() (messaging.Broadcaster, func()) +// Carrier is what a factory returns: two instances of one carrier over the same +// backing store, as two replicas of a deployment would have, and what the suite +// needs to know about the carrier. +type Carrier struct { + // Bus and Peer are two ends of the carrier. A message published on one must + // reach a subscriber on the other. A carrier that has no second end, such as + // an in-memory double, may return the same value twice. + Bus, Peer messaging.Broadcaster + // Cleanup releases both ends. + Cleanup func() + // ServesControlRoots is true when the carrier accepts the control roots + // (nodes, mcp), which carry request and reply traffic. A carrier that has no + // request and reply refuses them with messaging.ErrUnservedSubject. + ServesControlRoots bool + // Dropped returns how many messages the carrier has dropped since it was + // built because a consumer fell behind. It is nil for a carrier that does + // not count. + Dropped func() uint64 +} + +// Factory returns a ready carrier. A factory that cannot build one calls Skip or +// Fail itself. +type Factory func() Carrier + +// BroadcastRootSubjects names one subject for each broadcast root that every +// carrier serves. A spec in the messaging package asserts that the keys are the +// roots that messaging.ValidateSubject accepts as broadcast roots, so a root +// added there without a row here turns a spec red. +var BroadcastRootSubjects = map[string]string{ + "jobs": "jobs.j1.progress", + "agent": "agent.a1.events.u1", + "gallery": "gallery.g1.progress", + "cache": "cache.invalidate.models", + "staging": "staging.m1.progress", + "prefixcache": "prefixcache.observe", + "responses": "responses.r1.cancel", + "state": "state.s1", + "finetune": "finetune.f1.progress", +} + +// controlRootSubjects names one subject for each control root. +var controlRootSubjects = map[string]string{ + "nodes": "nodes.n1.backend.stop", + "mcp": "mcp.a1.tools", +} type blob struct { Data string `json:"data"` @@ -53,15 +96,19 @@ func RunBroadcasterConformance(newBus Factory) { Describe("Broadcaster conformance", func() { var ( bus messaging.Broadcaster - cleanup func() + peer messaging.Broadcaster + carrier Carrier ) - BeforeEach(func() { bus, cleanup = newBus() }) + BeforeEach(func() { + carrier = newBus() + bus, peer = carrier.Bus, carrier.Peer + }) // A factory that skips (no Docker for the NATS run) returns before - // setting cleanup, so the guard keeps a skip from turning into a panic. + // setting a cleanup, so the guard keeps a skip from turning into a panic. AfterEach(func() { - if cleanup != nil { - cleanup() + if carrier.Cleanup != nil { + carrier.Cleanup() } }) @@ -132,7 +179,7 @@ func RunBroadcasterConformance(newBus Factory) { Expect(errors.Is(err, messaging.ErrUnsupportedWildcard)).To(BeTrue(), "got %v", err) }) - DescribeTable("delivers payloads past the PostgreSQL notify cap intact", + DescribeTable("delivers payloads of every size around the PostgreSQL notify cap intact", func(size int) { c := &collector{} _, err := bus.Subscribe("cache.invalidate.models", c.handler) @@ -147,8 +194,99 @@ func RunBroadcasterConformance(newBus Factory) { Expect(got.Data).To(HaveLen(size)) Expect(got.Data).To(Equal(want)) }, - Entry("just under the notify cap", 7900), - Entry("over the notify cap", 64*1024), + Entry("one byte", 1), + Entry("the last size that fits in a notification", 7999), + Entry("the first size that does not", 8000), + Entry("past the cap", 64*1024), + ) + + It("refuses a payload above the bound that every carrier shares", func() { + c := &collector{} + _, err := bus.Subscribe("cache.invalidate.models", c.handler) + Expect(err).ToNot(HaveOccurred()) + + err = bus.Publish("cache.invalidate.models", blob{Data: strings.Repeat("a", messaging.MaxBroadcastBytes)}) + Expect(errors.Is(err, messaging.ErrPayloadTooLarge)).To(BeTrue(), "got %v", err) + Consistently(c.count, 300*time.Millisecond).Should(BeZero()) + }) + + DescribeTable("carries a message on every broadcast root", + func(root string) { + subject := BroadcastRootSubjects[root] + c := &collector{} + _, err := peer.Subscribe(subject, c.handler) + Expect(err).ToNot(HaveOccurred()) + + Expect(bus.Publish(subject, blob{Data: root})).To(Succeed()) + + Eventually(c.count, 5*time.Second).Should(Equal(1)) + }, + Entry("jobs", "jobs"), + Entry("agent", "agent"), + Entry("gallery", "gallery"), + Entry("cache", "cache"), + Entry("staging", "staging"), + Entry("prefixcache", "prefixcache"), + Entry("responses", "responses"), + Entry("state", "state"), + Entry("finetune", "finetune"), + ) + + It("carries a message to a subscriber on the other end of the carrier", func() { + here, there := &collector{}, &collector{} + _, err := bus.Subscribe("jobs.j2.result", here.handler) + Expect(err).ToNot(HaveOccurred()) + _, err = peer.Subscribe("jobs.j2.result", there.handler) + Expect(err).ToNot(HaveOccurred()) + + Expect(bus.Publish("jobs.j2.result", blob{Data: "x"})).To(Succeed()) + + Eventually(there.count, 5*time.Second).Should(Equal(1)) + Eventually(here.count, 5*time.Second).Should(Equal(1)) + }) + + It("carries a payload past the notify cap to the other end", func() { + c := &collector{} + _, err := peer.Subscribe("jobs.j3.result", c.handler) + Expect(err).ToNot(HaveOccurred()) + + want := strings.Repeat("b", 64*1024) + Expect(bus.Publish("jobs.j3.result", blob{Data: want})).To(Succeed()) + + Eventually(c.count, 10*time.Second).Should(Equal(1)) + var got blob + Expect(json.Unmarshal(c.first(), &got)).To(Succeed()) + Expect(got.Data).To(Equal(want)) + }) + + DescribeTable("handles the control roots as the carrier declares", + func(root string) { + subject := controlRootSubjects[root] + err := bus.Publish(subject, blob{Data: "x"}) + if carrier.ServesControlRoots { + Expect(err).ToNot(HaveOccurred()) + return + } + Expect(errors.Is(err, messaging.ErrUnservedSubject)).To(BeTrue(), "publish: %v", err) + _, err = bus.Subscribe(subject, func([]byte) {}) + Expect(errors.Is(err, messaging.ErrUnservedSubject)).To(BeTrue(), "subscribe: %v", err) + }, + Entry("nodes", "nodes"), + Entry("mcp", "mcp"), ) + + It("reports no dropped message when every consumer keeps up", func() { + if carrier.Dropped == nil { + Skip("the carrier does not count dropped messages") + } + c := &collector{} + _, err := bus.Subscribe("gallery.g2.progress", c.handler) + Expect(err).ToNot(HaveOccurred()) + + Expect(bus.Publish("gallery.g2.progress", blob{Data: "x"})).To(Succeed()) + + Eventually(c.count, 5*time.Second).Should(Equal(1)) + Expect(carrier.Dropped()).To(BeZero()) + }) }) } diff --git a/core/services/testutil/export_test.go b/core/services/testutil/export_test.go deleted file mode 100644 index cc6da89af7b6..000000000000 --- a/core/services/testutil/export_test.go +++ /dev/null @@ -1,4 +0,0 @@ -package testutil - -// SubjectMatches exposes the fake bus matching rule to the external specs. -var SubjectMatches = subjectMatches diff --git a/core/services/testutil/fakebus.go b/core/services/testutil/fakebus.go index 8d99d02fab03..607a57bd2d1e 100644 --- a/core/services/testutil/fakebus.go +++ b/core/services/testutil/fakebus.go @@ -2,7 +2,6 @@ package testutil import ( "encoding/json" - "strings" "sync" "time" @@ -69,12 +68,15 @@ func (b *FakeBus) Publish(subject string, data any) error { if err != nil { return err } + if err := messaging.CheckBroadcastSize(subject, len(payload)); err != nil { + return err + } b.mu.Lock() b.publishCounts[subject]++ subs := append([]fakeBusSub(nil), b.subs...) b.mu.Unlock() for _, s := range subs { - if subjectMatches(s.subject, subject) { + if messaging.SubjectMatches(s.subject, subject) { s.handler(payload) } } @@ -203,26 +205,3 @@ func (b *FakeBus) TriggerReconnect() { cb() } } - -// subjectMatches reports whether a subscription filter matches a concrete -// subject, honouring the single-token `*` wildcard the way NATS does, so the -// fake delivers to the same subscribers the real carrier would. -func subjectMatches(filter, subject string) bool { - if filter == subject { - return true - } - fp := strings.Split(filter, ".") - sp := strings.Split(subject, ".") - if len(fp) != len(sp) { - return false - } - for i := range fp { - if fp[i] == "*" { - continue - } - if fp[i] != sp[i] { - return false - } - } - return true -} diff --git a/core/services/testutil/fakebus_conformance_test.go b/core/services/testutil/fakebus_conformance_test.go index 05f40ebeceab..a52cefad2a13 100644 --- a/core/services/testutil/fakebus_conformance_test.go +++ b/core/services/testutil/fakebus_conformance_test.go @@ -3,13 +3,13 @@ package testutil_test import ( . "github.com/onsi/ginkgo/v2" - "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/messaging/messagingtest" "github.com/mudler/LocalAI/core/services/testutil" ) var _ = Describe("FakeBus", func() { - messagingtest.RunBroadcasterConformance(func() (messaging.Broadcaster, func()) { - return testutil.NewFakeBus(), func() {} + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { + bus := testutil.NewFakeBus() + return messagingtest.Carrier{Bus: bus, Peer: bus, ServesControlRoots: true} }) }) diff --git a/core/services/testutil/subject_match_test.go b/core/services/testutil/subject_match_test.go deleted file mode 100644 index 8ad5957b8d9c..000000000000 --- a/core/services/testutil/subject_match_test.go +++ /dev/null @@ -1,21 +0,0 @@ -package testutil_test - -import ( - . "github.com/onsi/ginkgo/v2" - . "github.com/onsi/gomega" - - "github.com/mudler/LocalAI/core/services/testutil" -) - -var _ = Describe("FakeBus subject matching", func() { - DescribeTable("matches like a NATS single-token wildcard", - func(filter, subject string, want bool) { - Expect(testutil.SubjectMatches(filter, subject)).To(Equal(want)) - }, - Entry("exact", "jobs.new", "jobs.new", true), - Entry("wildcard hit", "jobs.*.cancel", "jobs.abc.cancel", true), - Entry("wildcard wrong tail", "jobs.*.cancel", "jobs.abc.result", false), - Entry("wildcard does not span tokens", "jobs.*", "jobs.a.b", false), - Entry("length mismatch", "jobs.new", "jobs.new.extra", false), - ) -}) From c3c6af242a3dd97d8f4af5f434c2d752c772e501 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:27:38 +0000 Subject: [PATCH 012/162] test(testutil): keep the database in memory and return its connection string The postgres image declares its data directory as a volume. A removed container leaves that anonymous volume behind, so every spec that started a database left tens of megabytes on the disk of the host. Put the data directory on a tmpfs. Also return the connection string, for code that needs a connection of its own, such as a LISTEN connection that cannot come from a pool. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/testutil/testdb.go | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/core/services/testutil/testdb.go b/core/services/testutil/testdb.go index 80e511201b7d..3d1b4060f31e 100644 --- a/core/services/testutil/testdb.go +++ b/core/services/testutil/testdb.go @@ -19,6 +19,14 @@ import ( // SetupTestDB creates a fresh PostgreSQL 16 container and returns a gorm.DB. // The container is cleaned up via DeferCleanup when the test completes. func SetupTestDB() *gorm.DB { + db, _ := SetupTestDBWithDSN() + return db +} + +// SetupTestDBWithDSN is SetupTestDB that also returns the connection string of +// the container, for code that opens a connection of its own, such as a +// LISTEN connection that cannot come from a pool. +func SetupTestDBWithDSN() (*gorm.DB, string) { if runtime.GOOS == "darwin" { Skip("testcontainers requires Docker, not available on macOS CI") } @@ -27,6 +35,10 @@ func SetupTestDB() *gorm.DB { tcpostgres.WithDatabase("testdb"), tcpostgres.WithUsername("test"), tcpostgres.WithPassword("test"), + // The data directory lives in memory. The image declares it as a volume, + // and a container that is removed leaves that anonymous volume behind, + // so every spec would leave tens of megabytes on the disk of the host. + testcontainers.WithTmpfs(map[string]string{"/var/lib/postgresql/data": "rw"}), testcontainers.WithWaitStrategyAndDeadline(60*time.Second, wait.ForLog("database system is ready to accept connections").WithOccurrence(2)), ) @@ -38,5 +50,5 @@ func SetupTestDB() *gorm.DB { Logger: logger.Default.LogMode(logger.Silent), }) Expect(err).ToNot(HaveOccurred()) - return db + return db, connStr } From 05ff50148c5df89819eeb37bb60174e038e4db24 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:27:38 +0000 Subject: [PATCH 013/162] feat(messaging): refuse the control roots on a carrier with no request and reply A carrier that has fan-out and no request and reply serves the broadcast roots only. ValidateBroadcastSubject applies the shared rules and then refuses the nodes and mcp roots with the same class of error as a root that is not served. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/messaging/subject_rules.go | 15 +++++++++++ core/services/messaging/subject_rules_test.go | 26 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/core/services/messaging/subject_rules.go b/core/services/messaging/subject_rules.go index da0193bfe422..932c65aaf99b 100644 --- a/core/services/messaging/subject_rules.go +++ b/core/services/messaging/subject_rules.go @@ -62,3 +62,18 @@ func ValidateSubject(subject string) error { } return fmt.Errorf("%w: %q", ErrUnservedSubject, subject) } + +// ValidateBroadcastSubject is ValidateSubject for a carrier that has fan-out +// and no request and reply. It refuses the control roots with the same class +// of error as a root that is not served at all, because for that carrier they +// are not served. +func ValidateBroadcastSubject(subject string) error { + if err := ValidateSubject(subject); err != nil { + return err + } + root, _, _ := strings.Cut(subject, ".") + if _, ok := broadcastRoots[root]; !ok { + return fmt.Errorf("%w: %q carries request and reply, which this carrier does not", ErrUnservedSubject, subject) + } + return nil +} diff --git a/core/services/messaging/subject_rules_test.go b/core/services/messaging/subject_rules_test.go index 9f9ac0b64ed6..e90d4902ff17 100644 --- a/core/services/messaging/subject_rules_test.go +++ b/core/services/messaging/subject_rules_test.go @@ -98,3 +98,29 @@ var _ = Describe("Subject rules", func() { } }) }) + +var _ = Describe("ValidateBroadcastSubject", func() { + DescribeTable("accepts the subjects of the broadcast roots", + func(subject string) { + Expect(messaging.ValidateBroadcastSubject(subject)).To(Succeed()) + }, + Entry("exact", "jobs.j1.progress"), + Entry("wildcard", "agent.*.events.*"), + Entry("finetune", "finetune.job1.progress"), + ) + + DescribeTable("refuses a control root with the class of an unserved subject", + func(subject string) { + err := messaging.ValidateBroadcastSubject(subject) + Expect(errors.Is(err, messaging.ErrUnservedSubject)).To(BeTrue(), "got %v", err) + Expect(err.Error()).To(ContainSubstring(subject)) + }, + Entry("nodes", "nodes.n1.backend.stop"), + Entry("mcp", "mcp.a1.tools"), + ) + + It("keeps every refusal that ValidateSubject makes", func() { + Expect(errors.Is(messaging.ValidateBroadcastSubject("bogus.x"), messaging.ErrUnservedSubject)).To(BeTrue()) + Expect(errors.Is(messaging.ValidateBroadcastSubject("jobs.>"), messaging.ErrUnsupportedWildcard)).To(BeTrue()) + }) +}) From cccf81b4fb88a26843a39d5065dc595ff9da79e6 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:27:38 +0000 Subject: [PATCH 014/162] feat(pgbus): add a broadcast carrier over PostgreSQL LISTEN and NOTIFY The second implementation of messaging.Broadcaster. A message goes out with pg_notify. When its encoded notification is 8000 bytes or more, the payload is written to the bus_messages table and the notification carries the row id. Every replica deletes the rows that are older than the retention, by the clock of the database. Delivery is at-most-once, as with NATS. The listener copies each notification into a bounded queue and never blocks, so the server is not held up by a slow consumer. A lost connection is dialled again and every channel is listened again, then the reconnect callbacks run. The subject rules are the shared ones, and the carrier refuses the control roots. It checks the shared payload bound. Counters for the messages published by path, the messages lost by stage, the spilled rows deleted and the reconnections go to OpenTelemetry, because a lost broadcast is otherwise visible only in a log. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/pgbus/bus.go | 600 ++++++++++++++++++ core/services/pgbus/bus_test.go | 442 +++++++++++++ core/services/pgbus/conformance_test.go | 32 + core/services/pgbus/export_test.go | 23 + .../pgbus/listen_ordering_internal_test.go | 176 +++++ core/services/pgbus/listener.go | 352 ++++++++++ core/services/pgbus/listener_test.go | 312 +++++++++ core/services/pgbus/metrics.go | 84 +++ core/services/pgbus/metrics_test.go | 181 ++++++ core/services/pgbus/pgbus_suite_test.go | 13 + core/services/pgbus/spill.go | 49 ++ core/services/pgbus/spill_test.go | 272 ++++++++ go.mod | 2 +- 13 files changed, 2537 insertions(+), 1 deletion(-) create mode 100644 core/services/pgbus/bus.go create mode 100644 core/services/pgbus/bus_test.go create mode 100644 core/services/pgbus/conformance_test.go create mode 100644 core/services/pgbus/export_test.go create mode 100644 core/services/pgbus/listen_ordering_internal_test.go create mode 100644 core/services/pgbus/listener.go create mode 100644 core/services/pgbus/listener_test.go create mode 100644 core/services/pgbus/metrics.go create mode 100644 core/services/pgbus/metrics_test.go create mode 100644 core/services/pgbus/pgbus_suite_test.go create mode 100644 core/services/pgbus/spill.go create mode 100644 core/services/pgbus/spill_test.go diff --git a/core/services/pgbus/bus.go b/core/services/pgbus/bus.go new file mode 100644 index 000000000000..2a20ec147ad8 --- /dev/null +++ b/core/services/pgbus/bus.go @@ -0,0 +1,600 @@ +// Package pgbus carries fan-out broadcasts between frontend replicas over the +// PostgreSQL that a distributed deployment already runs, using LISTEN and +// NOTIFY. +// +// It exists so that a deployment needs a database and nothing else. It is the +// fan-out half of the NATS surface only: request and reply and queue groups are +// not here. +// +// Delivery is at-most-once, on purpose. A NOTIFY reaches the sessions that are +// listening when it is issued and nobody else, so a replica that is +// reconnecting misses what was published in that window, as a NATS subscriber +// does. Two consequences matter: +// +// - Nothing may read a message that did not arrive as evidence about a node. +// A carrier that cannot deliver is not a worker that is gone. Absence is +// read from the database and never from silence on a bus. +// - Anything that must survive a gap belongs in a table, and the broadcast is +// a hint to go and look. +// +// Limits, measured on a four-core database: about 16,000 broadcasts a second of +// 200 bytes, about 8,000 a second with 20 listeners or with 2 to 8 KiB, and +// about 1,000 a second of 64 KiB. A broadcast that does not fit in a +// notification is written to a row and costs a SELECT on every replica. +package pgbus + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/mudler/xlog" + "go.opentelemetry.io/otel/metric" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// channelPrefix namespaces the LISTEN channels of this carrier away from +// anything else that uses the same database. +const channelPrefix = "localai_" + +// maxNotifyPayloadBytes is the limit of PostgreSQL for a pg_notify payload, and +// it is exclusive: async.c refuses a payload when strlen(payload) >= +// NOTIFY_PAYLOAD_MAX_LENGTH. 7999 bytes is the largest payload that the server +// accepts, and 8000 fails. Checked against postgres:16. +// +// The size that is measured against it is the encoded notification and not the +// payload of the caller. The subject, the JSON escaping and the keys of the +// envelope travel too, so a payload that is sized against the limit alone +// would ship 7999 bytes of data in a notification that the server rejects. +const maxNotifyPayloadBytes = 8000 + +// deliveryQueueDepth is how far one subscriber may fall behind before its +// broadcasts are dropped. +// +// It is per subscriber and not per carrier. Handlers write to SSE streams and to +// channels of the caller, and running them on the listener goroutine would let +// one blocked writer stop delivery for every subject in the deployment. The drop +// is logged at error level and counted: to lose a broadcast loudly is +// recoverable, and to wedge the carrier silently is not. +const deliveryQueueDepth = 256 + +// ChannelFor maps a subject onto the LISTEN channel that carries it. All +// subjects of one root share one channel. +// +// Publish and Subscribe both call it, and neither may work out a channel name on +// its own: a publisher and a subscriber that disagree about a channel is a +// subject that is published and delivered to nobody, with no error anywhere. +// +// The channel is the root and not the subject because a channel name is limited +// to 63 bytes, and a channel for every subject would mean one LISTEN for every +// job, every gallery operation and every response, which has no bound. The set +// of roots is closed (see messaging.ValidateBroadcastSubject), so the set of +// channels is closed too. +func ChannelFor(subject string) (string, error) { + if err := messaging.ValidateBroadcastSubject(subject); err != nil { + return "", err + } + root, _, _ := strings.Cut(subject, ".") + return channelPrefix + root, nil +} + +// Config configures the carrier of a replica. +type Config struct { + // DSN is the connection string of the LISTEN connection. It is separate from + // DB because a LISTEN connection cannot come from a pool: a registration + // belongs to one backend session, so a pooled handle would register on one + // connection and lose it on the next. This connection is pinned for the life + // of the Bus. + // + // It comes from one place: config.ApplicationConfig.Auth.DatabaseURL. It is + // not a flag of its own and is not read from the environment. A second knob + // would let a deployment point the LISTEN connection and the pool at + // different databases, which is a carrier that publishes and never delivers. + // New refuses that combination and does not rely on the convention. + DSN string + // DB is the pooled handle that NOTIFY and the spill table are written on. + DB *gorm.DB + // Queue bounds how many notifications may wait for the resolver before this + // carrier drops them. Zero means DefaultQueueDepth. + Queue int + // Retention is how long a spilled row outlives its notification. Zero means + // DefaultSpillRetention. + Retention time.Duration + // SweepInterval is how often this carrier deletes the spilled rows that aged + // out. Zero means Retention / 2. + SweepInterval time.Duration + // Meter receives the counters of the carrier. Nil means the global meter. + Meter metric.Meter +} + +// notification is what travels in a pg_notify payload. The keys are one byte +// long because every byte of the envelope is a byte that the payload of the +// caller cannot use before it has to spill. +type notification struct { + Subject string `json:"s"` + Data json.RawMessage `json:"d,omitempty"` + SpillID string `json:"i,omitempty"` +} + +// Bus is one replica's end of the carrier. +type Bus struct { + cfg Config + ctx context.Context + cancel context.CancelFunc + + // connected is a fact about the link of this process to the database and + // about nothing else. Nothing may decide on it. "The carrier is down" is not + // one of the conditions a node can be in, and to read it as absence would let + // a database error evict models. + connected atomic.Bool + + // dropped counts the notifications that the listener or the resolver lost. + // Like connected, it may be reported, and nothing may act on it. + dropped atomic.Uint64 + + // appName identifies the LISTEN session of this carrier in pg_stat_activity, + // and retention is Config.Retention after the default. New sets both once and + // other goroutines read them, so neither may be assigned again. + appName string + retention time.Duration + metrics *metrics + + cmds chan listenCmd + inbound chan inbound + listenerDone chan struct{} + resolverDone chan struct{} + purgerDone chan struct{} + closeOnce sync.Once + + // cbMu guards reconnectCbs. It is a lock of its own because a registration + // comes from the goroutine of a caller while the listener dials again, and + // neither should wait for delivery to finish. + cbMu sync.Mutex + reconnectCbs []func() + + // listenMu orders the decision on the count of a channel with the LISTEN or + // UNLISTEN that the decision needs, as one step. + // + // To decide under mu and to issue outside it is a real defect. The last + // Unsubscribe may decide to UNLISTEN, a Subscribe may then decide to LISTEN, + // and the two may reach the connection in that order: the channel ends up not + // listened, with a live subscription on it. This does not heal itself, + // because the next Subscribe sees that it is not the first and never listens + // again, so the whole root is deaf on this replica until a connection loss + // makes the carrier listen again. + // + // It is not mu because command waits for the listener goroutine and delivery + // takes mu, so to hold mu across that wait would deadlock the carrier. The + // listener and the resolver never take this lock. + listenMu sync.Mutex + + // listenBarrier is a seam for tests and is nil in production. See barrier. + listenBarrier func(stage, op string) + + mu sync.Mutex + nextID uint64 + subs map[string]map[uint64]*subscription +} + +// New opens the carrier: one pinned LISTEN connection, and the pooled handle +// that publishes travel on. +func New(ctx context.Context, cfg Config) (*Bus, error) { + if cfg.DB == nil { + return nil, errors.New("pgbus: no database handle") + } + // Everything below is PostgreSQL only (pg_notify, LISTEN, now(), + // make_interval). Without this check, now() on SQLite reads as a missing + // migration and not as a carrier that was never meant to run there. + if name := cfg.DB.Dialector.Name(); name != "postgres" { + return nil, fmt.Errorf("pgbus: the broadcast carrier requires PostgreSQL, got %q", name) + } + if cfg.DSN == "" { + return nil, errors.New("pgbus: no DSN for the LISTEN connection") + } + + // The name is per carrier and not per deployment, so two replicas on one + // database are different in pg_stat_activity, and a spec can drop the session + // of one without touching the peer that it asserts against. + appName := applicationNamePrefix + uuid.NewString() + + conn, err := dial(ctx, cfg.DSN, appName) + if err != nil { + return nil, fmt.Errorf("pgbus: opening the LISTEN connection: %w", err) + } + if err := verifySameDatabase(ctx, conn, cfg.DB); err != nil { + _ = conn.Close(ctx) + return nil, err + } + + queue := cfg.Queue + if queue <= 0 { + queue = DefaultQueueDepth + } + retention := cfg.Retention + if retention <= 0 { + retention = DefaultSpillRetention + } + + busCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) + b := &Bus{ + cfg: cfg, + ctx: busCtx, + cancel: cancel, + appName: appName, + retention: retention, + metrics: newMetrics(cfg.Meter), + cmds: make(chan listenCmd), + inbound: make(chan inbound, queue), + listenerDone: make(chan struct{}), + resolverDone: make(chan struct{}), + purgerDone: make(chan struct{}), + subs: map[string]map[uint64]*subscription{}, + } + b.connected.Store(true) + go b.listen(conn) + go b.resolve() + go b.purge() + return b, nil +} + +// verifySameDatabase refuses a carrier whose LISTEN connection and whose pool are +// not on the same database. +// +// That combination has no other symptom: every publish succeeds, every subscribe +// succeeds, and nothing is ever delivered, on every replica, for ever. It is +// cheap to exclude here and expensive to diagnose anywhere else. +func verifySameDatabase(ctx context.Context, conn *pgx.Conn, db *gorm.DB) error { + const identity = `SELECT current_database() || '@' || coalesce(host(inet_server_addr()), 'local') || ':' || coalesce(inet_server_port(), 0)` + + var listenSide string + if err := conn.QueryRow(ctx, identity).Scan(&listenSide); err != nil { + return fmt.Errorf("pgbus: identifying the database of the LISTEN connection: %w", err) + } + var poolSide string + if err := db.WithContext(ctx).Raw(identity).Row().Scan(&poolSide); err != nil { + return fmt.Errorf("pgbus: identifying the database of the pool: %w", err) + } + if listenSide != poolSide { + return fmt.Errorf("pgbus: the LISTEN connection and the pool are on a different database (%s vs %s); every broadcast would be published and never delivered", listenSide, poolSide) + } + return nil +} + +// IsConnected reports whether the pinned LISTEN connection is up. See the +// comment on Bus.connected: it is for reports and logs and not for decisions. +func (b *Bus) IsConnected() bool { return b.connected.Load() } + +// Close stops the carrier. It is idempotent and does not wait for handlers that +// are still running: a handler that never returns must not hold up the shutdown +// of its process. +func (b *Bus) Close() { + b.closeOnce.Do(func() { + b.cancel() + <-b.listenerDone + <-b.resolverDone + <-b.purgerDone + b.connected.Store(false) + }) +} + +// inlineNotification encodes what Publish puts on the wire when a message fits in +// the notification, and returns it with the encoded payload of the caller so +// that the spill path does not encode a second time. +func inlineNotification(subject string, data any) ([]byte, json.RawMessage, error) { + payload, err := json.Marshal(data) + if err != nil { + return nil, nil, fmt.Errorf("pgbus: encoding a broadcast on %q: %w", subject, err) + } + encoded, err := json.Marshal(notification{Subject: subject, Data: payload}) + if err != nil { + return nil, nil, fmt.Errorf("pgbus: encoding the notification for %q: %w", subject, err) + } + return encoded, payload, nil +} + +// Publish fans a message out to every subscriber of the subject on every +// replica, this one included. +func (b *Bus) Publish(subject string, data any) error { + channel, err := ChannelFor(subject) + if err != nil { + return err + } + encoded, payload, err := inlineNotification(subject, data) + if err != nil { + return err + } + if err := messaging.CheckBroadcastSize(subject, len(payload)); err != nil { + return err + } + // This is the one size decision in the package. Several subjects exceed the + // limit in normal operation (the result of a job carries the whole output of + // a model, and a gallery progress event carries one entry for each node), so + // the spill is the usual path for them and not an error. + path := pathInline + if len(encoded) >= maxNotifyPayloadBytes { + id, err := b.spill(subject, payload) + if err != nil { + return err + } + encoded, err = json.Marshal(notification{Subject: subject, SpillID: id}) + if err != nil { + return fmt.Errorf("pgbus: encoding the spill notification for %q: %w", subject, err) + } + path = pathSpill + } + + // pg_notify and not a NOTIFY statement, because the channel is a value here + // and NOTIFY would need it as an identifier in the text. + if err := b.cfg.DB.WithContext(b.ctx).Exec("SELECT pg_notify(?, ?)", channel, string(encoded)).Error; err != nil { + return fmt.Errorf("pgbus: publishing on %q: %w", subject, err) + } + b.metrics.publish(path) + return nil +} + +// Subscribe registers a handler for every subject that matches the filter. The +// grammar of the filter is the shared one: an exact subject, or `*` for a whole +// token. +func (b *Bus) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { + channel, err := ChannelFor(subject) + if err != nil { + return nil, err + } + if handler == nil { + return nil, fmt.Errorf("pgbus: no handler for %q", subject) + } + + // Before the lock, so a spec can see that this registration has started even + // when the lock is what stops it going further. + b.barrier("enter", "LISTEN") + b.listenMu.Lock() + defer b.listenMu.Unlock() + + b.mu.Lock() + b.nextID++ + sub := &subscription{ + bus: b, + channel: channel, + filter: subject, + id: b.nextID, + queue: make(chan []byte, deliveryQueueDepth), + stop: make(chan struct{}), + done: make(chan struct{}), + } + first := len(b.subs[channel]) == 0 + b.mu.Unlock() + + // The LISTEN comes before the registration and not after it. + // + // Both orders have listened by the time Subscribe returns, so the caller + // cannot tell them apart. An observer can. With the registration first, the + // count of subscribers includes a handler whose channel is not listened yet, + // so whoever waits on that count as a sign that the carrier is ready goes on + // into a window where a publish is lost. Registering last makes the count + // mean what its name says, and it removes the rollback of a failed LISTEN: + // a registration that was never made needs no undoing. + if first { + b.barrier("issue", "LISTEN") + if err := b.command("LISTEN " + pgx.Identifier{channel}.Sanitize()); err != nil { + close(sub.stop) + return nil, err + } + } + + b.mu.Lock() + if b.subs[channel] == nil { + b.subs[channel] = map[uint64]*subscription{} + } + b.subs[channel][sub.id] = sub + b.mu.Unlock() + + go sub.run(handler) + return sub, nil +} + +// subscribers reports how many handlers are registered, in every channel. A +// handler that is counted is live: its channel was listened before it was +// counted. +// +// The count exists to be asserted, because the leak that it shows has no other +// symptom. Two subscriptions in a deployment are opened and closed for every +// HTTP request (the progress stream of a job and the event stream of an agent), +// and only the first subscriber of a channel issues a LISTEN. The others only +// register a filter in the process. An Unsubscribe that did not remove its filter +// would leave a replica that served ten thousand requests running ten thousand +// closures for every notification, and nothing would fail. It would be slower. +func (b *Bus) subscribers() int { + b.mu.Lock() + defer b.mu.Unlock() + n := 0 + for _, channel := range b.subs { + n += len(channel) + } + return n +} + +// barrier is a seam for tests and does nothing in production. +// +// It exists because the order that listenMu enforces cannot be seen from outside +// the package and cannot be provoked from outside either. The natural window is +// microseconds wide: it had no hit in forty attempts, and ten out of ten once it +// was widened. A spec that waits for that window is a spec that passes by luck. +func (b *Bus) barrier(stage, op string) { + if b.listenBarrier != nil { + b.listenBarrier(stage, op) + } +} + +// deliver resolves one notification and hands it to the subscribers whose +// filters match. +func (b *Bus) deliver(channel, payload string) { + b.barrier("enter", "DELIVER") + + var n notification + if err := json.Unmarshal([]byte(payload), &n); err != nil { + b.lose(stageResolve) + xlog.Error("Broadcast carrier received an undecodable notification", "channel", channel, "error", err) + return + } + + data := []byte(n.Data) + if n.SpillID == "" && len(data) == 0 { + // Publish cannot produce this, because json.Marshal never returns an + // empty encoding. It is possible only if something other than this + // carrier notifies on one of its channels. A handler that received nil + // would see a message that says nothing, and no message is better. + b.lose(stageResolve) + xlog.Error("Broadcast carrier received a notification with no payload", "channel", channel, "subject", n.Subject) + return + } + if n.SpillID != "" { + resolved, err := b.resolveSpill(n.SpillID) + if err != nil { + // A dropped broadcast and nothing more. A handler that got an empty + // message instead would read a carrier failure as a fact about the + // deployment. + b.lose(stageResolve) + xlog.Error("Broadcast carrier could not resolve a spilled message", "subject", n.Subject, "id", n.SpillID, "error", err) + return + } + data = resolved + } + + b.mu.Lock() + targets := make([]*subscription, 0, len(b.subs[channel])) + for _, sub := range b.subs[channel] { + // The shared matcher decides here and nowhere else. A second spelling is + // the drift that the shared function exists to prevent. + if messaging.SubjectMatches(sub.filter, n.Subject) { + targets = append(targets, sub) + } + } + b.mu.Unlock() + + for _, sub := range targets { + sub.enqueue(n.Subject, data) + } +} + +// lose counts a broadcast that this replica received and did not deliver. +func (b *Bus) lose(stage string) { + b.dropped.Add(1) + b.metrics.drop(stage) +} + +// subscription is the registration of one handler. +type subscription struct { + bus *Bus + channel string + filter string + id uint64 + + queue chan []byte + stop chan struct{} + done chan struct{} + dropped atomic.Uint64 + once sync.Once +} + +// DropCounter is the optional interface of a Subscription that can report how +// many broadcasts it lost. +// +// It exists because the loss is otherwise invisible to the party that needs to +// know. The error log is on the replica that received the broadcast, and there is +// no sequence number and no signal of a gap, so a subscriber cannot act on "a +// broadcast that must survive a gap belongs in a table". A subscriber whose +// subject has no next message, such as the result of a job, can ask for this +// interface and go to read the row. +// +// It is not on messaging.Subscription. A subscription that only this carrier can +// answer for does not belong on the interface of every carrier, and a wider +// interface would make every consumer claim a guarantee that it does not have. +type DropCounter interface { + Dropped() uint64 +} + +// Dropped reports how many broadcasts this subscription lost because its handler +// was too far behind. It only grows. +func (s *subscription) Dropped() uint64 { return s.dropped.Load() } + +func (s *subscription) run(handler func([]byte)) { + defer close(s.done) + for { + select { + case data := <-s.queue: + handler(data) + case <-s.stop: + return + case <-s.bus.ctx.Done(): + return + } + } +} + +func (s *subscription) enqueue(subject string, data []byte) { + select { + case s.queue <- data: + default: + s.dropped.Add(1) + s.bus.lose(stageSubscription) + xlog.Error("Broadcast carrier dropped a message: subscriber is not keeping up", + "filter", s.filter, "subject", subject, "depth", deliveryQueueDepth, "dropped", s.dropped.Load()) + } +} + +// Unsubscribe stops delivery to this handler and, when it was the last +// subscriber of its channel, stops listening on the channel. +func (s *subscription) Unsubscribe() error { + var err error + s.once.Do(func() { + b := s.bus + b.barrier("enter", "UNLISTEN") + // The whole decision and its issue, as one step. See listenMu. + b.listenMu.Lock() + defer b.listenMu.Unlock() + + b.mu.Lock() + delete(b.subs[s.channel], s.id) + last := len(b.subs[s.channel]) == 0 + if last { + delete(b.subs, s.channel) + } + b.mu.Unlock() + + close(s.stop) + if last && b.ctx.Err() == nil { + b.barrier("issue", "UNLISTEN") + err = b.command("UNLISTEN " + pgx.Identifier{s.channel}.Sanitize()) + } + }) + return err +} + +// spill writes a broadcast that does not fit in a notification to a row, and +// returns the id that the notification carries in its place. +func (b *Bus) spill(subject string, payload []byte) (string, error) { + row := BusMessage{ID: uuid.New().String(), Subject: subject, Payload: payload} + if err := b.cfg.DB.WithContext(b.ctx).Create(&row).Error; err != nil { + return "", fmt.Errorf("pgbus: spilling a broadcast on %q: %w", subject, err) + } + return row.ID, nil +} + +func (b *Bus) resolveSpill(id string) ([]byte, error) { + var row BusMessage + if err := b.cfg.DB.WithContext(b.ctx).Where("id = ?", id).First(&row).Error; err != nil { + return nil, err + } + return row.Payload, nil +} + +var _ messaging.Broadcaster = (*Bus)(nil) diff --git a/core/services/pgbus/bus_test.go b/core/services/pgbus/bus_test.go new file mode 100644 index 000000000000..217b19e68893 --- /dev/null +++ b/core/services/pgbus/bus_test.go @@ -0,0 +1,442 @@ +package pgbus_test + +import ( + "context" + "errors" + "path/filepath" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// Every delivery spec here builds TWO buses on TWO connections and asserts on +// the one that did NOT publish. One bus that hears itself proves nothing about +// a carrier whose whole job is to reach the OTHER replica, and an in-memory +// double proves less than that: it cannot fail the way a LISTEN connection +// fails. +var _ = Describe("the PostgreSQL broadcast carrier", func() { + var ( + db *gorm.DB + dsn string + pub, sub *pgbus.Bus + ) + + // newBus is the only construction path in this file, so a spec cannot + // accidentally skip the migration and read a missing-table error as a + // delivery failure. + 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 + } + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + pub = newBus() + sub = newBus() + }) + + // received registers a subscriber that forwards every delivery onto a + // channel, so specs synchronize on delivery rather than on the clock. + received := func(b *pgbus.Bus, filter string) (<-chan []byte, messaging.Subscription) { + GinkgoHelper() + out := make(chan []byte, 16) + s, err := b.Subscribe(filter, func(data []byte) { out <- data }) + Expect(err).ToNot(HaveOccurred()) + Expect(s).ToNot(BeNil()) + return out, s + } + + Describe("delivery across two connections", func() { + It("carries a message to a subscriber on the other bus", func() { + out, _ := received(sub, "jobs.abc.progress") + + Expect(pub.Publish("jobs.abc.progress", map[string]string{"phase": "downloading"})).To(Succeed()) + + Eventually(out).Should(Receive(MatchJSON(`{"phase":"downloading"}`))) + }) + + // The root set is not a formality: a subject family that maps to its own + // channel is a family a subscriber on a different channel never hears. + // A Publish or a Subscribe that hardcodes one channel shows up here as + // silence. + It("carries a message on a root other than jobs", func() { + out, _ := received(sub, "gallery.install.progress") + + Expect(pub.Publish("gallery.install.progress", map[string]int{"done": 3})).To(Succeed()) + + Eventually(out).Should(Receive(MatchJSON(`{"done":3}`))) + }) + + It("delivers to every subscriber of a subject, on the publishing bus and on the other one", func() { + // Fan-out is the whole point of this carrier: a replica that + // published must still run its own handlers, because in production + // the publisher is one of the replicas that has to react. + here, _ := received(pub, "cache.evict") + there, _ := received(sub, "cache.evict") + + Expect(pub.Publish("cache.evict", map[string]string{"model": "m"})).To(Succeed()) + + Eventually(here).Should(Receive()) + Eventually(there).Should(Receive()) + }) + }) + + Describe("which subjects reach which filters", func() { + // The expectations are COMPUTED by messaging.SubjectMatches rather than + // written out, so the carrier and the matcher cannot drift apart: if + // the carrier ever answers differently from the shared function, this + // table reddens without anyone having to notice the divergence. + // + // Ordering carries the negative half. Every case publishes the probe + // subject and then a sentinel on a subject the filter certainly + // matches; a filter that must NOT see the probe is proved by the + // sentinel arriving first, which needs no clock. + DescribeTable("agrees with messaging.SubjectMatches", + func(filter, subject string) { + out, _ := received(sub, filter) + sentinel := "jobs.sentinel.progress" + if !messaging.SubjectMatches(filter, sentinel) { + sentinel = filter + } + + Expect(pub.Publish(subject, map[string]string{"m": "probe"})).To(Succeed()) + Expect(pub.Publish(sentinel, map[string]string{"m": "sentinel"})).To(Succeed()) + + var first []byte + Eventually(out).Should(Receive(&first)) + want := `{"m":"sentinel"}` + if messaging.SubjectMatches(filter, subject) { + want = `{"m":"probe"}` + } + Expect(first).To(MatchJSON(want)) + }, + Entry("exact subject", "jobs.abc.progress", "jobs.abc.progress"), + Entry("single-token wildcard in the middle", "jobs.*.progress", "jobs.abc.progress"), + Entry("a sibling subject of the same family", "jobs.*.progress", "jobs.abc.result"), + Entry("more tokens than the filter has", "jobs.*.progress", "jobs.abc.extra.progress"), + Entry("fewer tokens than the filter has", "jobs.*.progress", "jobs.progress"), + Entry("a wildcard in the leading token", "jobs.*", "jobs.abc"), + ) + }) + + Describe("subjects this carrier refuses", func() { + It("refuses a publish whose root it does not serve, and names the subject", func() { + err := pub.Publish("weather.today", map[string]string{"sky": "blue"}) + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("weather.today")) + }) + + It("delivers nothing at all for a refused publish", func() { + // The refusal has to be a refusal, not an error returned beside a + // message that went somewhere anyway. + out, _ := received(sub, "jobs.abc.progress") + + Expect(pub.Publish("weather.today", map[string]string{"sky": "blue"})).ToNot(Succeed()) + Expect(pub.Publish("jobs.abc.progress", map[string]string{"m": "sentinel"})).To(Succeed()) + + var first []byte + Eventually(out).Should(Receive(&first)) + Expect(first).To(MatchJSON(`{"m":"sentinel"}`)) + }) + + It("refuses a subscribe whose root it does not serve, and registers nothing", func() { + s, err := sub.Subscribe("weather.today", func([]byte) {}) + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("weather.today")) + Expect(s).To(BeNil()) + }) + + It("refuses a filter the shared rules do not allow, with the error of that package", func() { + // errors.Is and not a string: this proves that the carrier asks + // messaging and does not spell the rule again. + s, err := sub.Subscribe("jobs.>", func([]byte) {}) + + Expect(errors.Is(err, messaging.ErrUnsupportedWildcard)).To(BeTrue(), "got %v", err) + Expect(s).To(BeNil()) + }) + + It("refuses an empty-token filter with the class of an unserved subject", func() { + s, err := sub.Subscribe("jobs..progress", func([]byte) {}) + + Expect(errors.Is(err, messaging.ErrUnservedSubject)).To(BeTrue(), "got %v", err) + Expect(s).To(BeNil()) + }) + + It("refuses the control roots, which carry request and reply", func() { + err := pub.Publish("nodes.n1.backend.stop", map[string]string{"m": "x"}) + Expect(errors.Is(err, messaging.ErrUnservedSubject)).To(BeTrue(), "got %v", err) + + s, err := sub.Subscribe("mcp.a1.tools", func([]byte) {}) + Expect(errors.Is(err, messaging.ErrUnservedSubject)).To(BeTrue(), "got %v", err) + Expect(s).To(BeNil()) + }) + }) + + Describe("unsubscribing", func() { + It("stops delivery to that handler and leaves the others alone", func() { + gone, s := received(sub, "jobs.abc.progress") + stays, _ := received(sub, "jobs.abc.progress") + + Expect(s.Unsubscribe()).To(Succeed()) + Expect(pub.Publish("jobs.abc.progress", map[string]string{"m": "after"})).To(Succeed()) + + Eventually(stays).Should(Receive(MatchJSON(`{"m":"after"}`))) + Consistently(gone).ShouldNot(Receive()) + }) + }) + + Describe("subscriptions that come and go", func() { + // The accumulation assertion, and it is deliberately NOT a delivery + // one. Two subscriptions in this deployment are opened and closed per + // HTTP REQUEST, so a filter that outlives its Unsubscribe is a replica + // that gets steadily slower and never fails. Delivery specs cannot see + // it: the leaked handlers deliver correctly, there are just more of + // them every request. + It("registers nothing that outlives its Unsubscribe", func() { + delivered := make(chan []byte, 8) + survivor, err := sub.Subscribe("jobs.churn.progress", func(data []byte) { delivered <- data }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = survivor.Unsubscribe() }) + + start := sub.Subscribers() + Expect(start).To(Equal(1)) + + for i := 0; i < 1000; i++ { + s, err := sub.Subscribe("jobs.churn.progress", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(s.Unsubscribe()).To(Succeed()) + } + + Expect(sub.Subscribers()).To(Equal(start), + "a subscription that is opened and closed must leave the register where it found it") + + // And the survivor is still exactly one subscriber, not zero and + // not a thousand. A counter rather than a boolean, because + // "something arrived" is true for every wrong answer here. + Expect(pub.Publish("jobs.churn.progress", map[string]string{"m": "after"})).To(Succeed()) + Eventually(delivered).Should(Receive(MatchJSON(`{"m":"after"}`))) + Consistently(delivered, "500ms", "50ms").ShouldNot(Receive(), + "a leaked handler would deliver the same notification again") + }) + + It("keeps the channel listened while a subscriber remains", func() { + // The churn above cycles LISTEN and UNLISTEN a thousand times on a + // channel that still has a live subscriber. If the refcount and the + // UNLISTEN it implies were not decided as one step, the channel + // ends up unlistened with a live subscription on it, and it does + // not self-heal: the next Subscribe sees first == false and never + // re-LISTENs. + out, _ := received(sub, "jobs.abc.progress") + for i := 0; i < 50; i++ { + s, err := sub.Subscribe("jobs.other.progress", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(s.Unsubscribe()).To(Succeed()) + } + + Expect(pub.Publish("jobs.abc.progress", map[string]string{"m": "still here"})).To(Succeed()) + Eventually(out).Should(Receive(MatchJSON(`{"m":"still here"}`))) + }) + }) + + Describe("what Subscribe has finished doing when it returns", func() { + // The per-request subscriptions are why this matters. An SSE client + // connects and then triggers the work; if Subscribe returns before the + // LISTEN is on the wire, the first events go nowhere and the stream + // looks hung with no error on either side. + // + // There is deliberately no Eventually around the PUBLISH. Retrying the + // publish is exactly the thing that would hide an asynchronous LISTEN. + It("has already issued the LISTEN, every time", func() { + for i := 0; i < 50; i++ { + out := make(chan []byte, 1) + s, err := sub.Subscribe("jobs.race.progress", func(data []byte) { out <- data }) + Expect(err).ToNot(HaveOccurred()) + + Expect(pub.Publish("jobs.race.progress", map[string]int{"i": i})).To(Succeed()) + + Eventually(out, "10s").Should(Receive(), "iteration %d subscribed and then missed its own first message", i) + Expect(s.Unsubscribe()).To(Succeed()) + } + }) + }) + + Describe("a handler that does not return", func() { + It("does not stop delivery to the other subscribers", func() { + // A carrier whose handlers run on one goroutine is a carrier one + // blocked SSE writer can wedge for the whole deployment. + block := make(chan struct{}) + DeferCleanup(func() { close(block) }) + _, _ = sub.Subscribe("jobs.abc.progress", func([]byte) { <-block }) + free, _ := received(sub, "jobs.abc.progress") + + Expect(pub.Publish("jobs.abc.progress", map[string]string{"m": "one"})).To(Succeed()) + Expect(pub.Publish("jobs.abc.progress", map[string]string{"m": "two"})).To(Succeed()) + + Eventually(free).Should(Receive(MatchJSON(`{"m":"one"}`))) + Eventually(free).Should(Receive(MatchJSON(`{"m":"two"}`))) + }) + }) + + Describe("a subscriber that falls behind", func() { + It("counts what it lost, so the subscriber can go and read the row", func() { + // The drop is loud in the log of the replica that took it, and that + // is the wrong party: the subscriber is the one that has to decide + // to go and read the table instead. A job result has no next + // message, so an invisible loss there is a job whose answer never + // arrives. + block := make(chan struct{}) + DeferCleanup(func() { close(block) }) + s, err := sub.Subscribe("jobs.slow.consumer", func([]byte) { <-block }) + Expect(err).ToNot(HaveOccurred()) + + counter, ok := s.(pgbus.DropCounter) + Expect(ok).To(BeTrue(), "a subscription must be able to report its losses") + Expect(counter.Dropped()).To(BeZero()) + + // Comfortably past the queue depth, stated absolutely rather than + // as deliveryQueueDepth+n, so a change to the depth moves the + // behaviour and not the expectation. + for i := 0; i < 400; i++ { + Expect(pub.Publish("jobs.slow.consumer", map[string]int{"i": i})).To(Succeed()) + } + + Eventually(counter.Dropped, 30*time.Second).Should(BeNumerically(">", 0)) + }) + }) + + Describe("losing the LISTEN connection", func() { + It("reconnects, restores its registrations and delivers again", func() { + // The transport failure a fake cannot produce. Terminating the + // backend sessions drops both buses' pinned connections; what must + // survive is the SUBSCRIPTION, because a carrier that comes back + // deaf looks exactly like a peer that stopped publishing. + out, _ := received(sub, "jobs.abc.progress") + + Expect(db.Exec( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid()", + ).Error).To(Succeed()) + + Eventually(sub.IsConnected, 30*time.Second, 100*time.Millisecond).Should(BeTrue()) + Eventually(func() bool { + // Retried because the publishing side's pool is recovering too, + // and because a NOTIFY that lands between the drop and the + // re-LISTEN is genuinely lost: this carrier is at-most-once. + _ = pub.Publish("jobs.abc.progress", map[string]string{"m": "after"}) + select { + case <-out: + return true + default: + return false + } + }, 30*time.Second, 200*time.Millisecond).Should(BeTrue()) + }) + }) +}) + +var _ = Describe("mapping a subject onto a LISTEN channel", func() { + It("puts a subject on the channel of its first token", func() { + Expect(pgbus.ChannelFor("jobs.abc.progress")).To(Equal("localai_jobs")) + }) + + It("serves finetune, which the shared rules list as a broadcast root", func() { + Expect(pgbus.ChannelFor("finetune.job1.progress")).To(Equal("localai_finetune")) + }) + + It("refuses a subject whose root it does not serve", func() { + _, err := pgbus.ChannelFor("weather.today") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("weather.today")) + }) + + It("refuses an empty subject", func() { + _, err := pgbus.ChannelFor("") + Expect(err).To(HaveOccurred()) + }) + + It("keeps every channel it can ever return inside PostgreSQL's identifier limit", func() { + // Asserted over the whole root set rather than over a sample, so a root + // added later cannot push a channel past 63 bytes silently. Past that + // limit PostgreSQL truncates the name instead of failing, which merges + // two families onto one channel. + Expect(messagingtest.BroadcastRootSubjects).ToNot(BeEmpty()) + for root := range messagingtest.BroadcastRootSubjects { + channel, err := pgbus.ChannelFor(root + ".anything") + Expect(err).ToNot(HaveOccurred(), root) + Expect(len(channel)).To(BeNumerically("<=", 63), channel) + } + }) +}) + +var _ = Describe("constructing the carrier", func() { + It("refuses a non-PostgreSQL handle and names the dialect", func() { + db, err := gorm.Open(sqlite.Open(filepath.Join(GinkgoT().TempDir(), "bus.db")), &gorm.Config{}) + Expect(err).ToNot(HaveOccurred()) + + _, err = pgbus.New(context.Background(), pgbus.Config{DSN: "postgres://x/y", DB: db}) + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("sqlite")) + }) + + It("refuses a DSN that points at a different database from the pool", func() { + // The failure this excludes has no other symptom: every publish + // succeeds, every subscribe succeeds, and nothing is ever delivered. + db, _ := testutil.SetupTestDBWithDSN() + _, otherDSN := testutil.SetupTestDBWithDSN() + + _, err := pgbus.New(context.Background(), pgbus.Config{DSN: otherDSN, DB: db}) + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("different database")) + }) + + It("reports itself connected while open and disconnected once closed", func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + + b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db}) + Expect(err).ToNot(HaveOccurred()) + Expect(b.IsConnected()).To(BeTrue()) + Expect(b.DSN()).To(Equal(dsn)) + + b.Close() + Expect(b.IsConnected()).To(BeFalse()) + }) + + It("survives a second Close", func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db}) + Expect(err).ToNot(HaveOccurred()) + + b.Close() + Expect(b.Close).ToNot(Panic()) + }) + + It("refuses an empty DSN rather than opening a carrier that cannot listen", func() { + db, _ := testutil.SetupTestDBWithDSN() + + _, err := pgbus.New(context.Background(), pgbus.Config{DB: db}) + + Expect(err).To(HaveOccurred()) + }) + + It("refuses a nil handle", func() { + _, err := pgbus.New(context.Background(), pgbus.Config{DSN: "postgres://x/y"}) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/pgbus/conformance_test.go b/core/services/pgbus/conformance_test.go new file mode 100644 index 000000000000..f36be720bb95 --- /dev/null +++ b/core/services/pgbus/conformance_test.go @@ -0,0 +1,32 @@ +package pgbus_test + +import ( + "context" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// The carrier runs the suite that every fan-out carrier runs, on two ends over +// one database as two replicas have. It serves the broadcast roots only. +var _ = Describe("the PostgreSQL broadcast carrier", func() { + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + open := func() *pgbus.Bus { + b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db}) + Expect(err).ToNot(HaveOccurred()) + return b + } + bus, peer := open(), open() + return messagingtest.Carrier{ + Bus: bus, Peer: peer, + Cleanup: func() { bus.Close(); peer.Close() }, + Dropped: func() uint64 { return bus.Dropped() + peer.Dropped() }, + } + }) +}) diff --git a/core/services/pgbus/export_test.go b/core/services/pgbus/export_test.go new file mode 100644 index 000000000000..d9910fadbb97 --- /dev/null +++ b/core/services/pgbus/export_test.go @@ -0,0 +1,23 @@ +package pgbus + +import ( + "context" + "time" +) + +// These expose internals to the specs. The production code has no caller for +// them. + +const ( + MaxNotifyPayloadBytes = maxNotifyPayloadBytes + SpillSweepSQL = spillSweepSQL +) + +func (b *Bus) Subscribers() int { return b.subscribers() } +func (b *Bus) ApplicationName() string { return b.appName } +func (b *Bus) QueueDepth() int { return cap(b.inbound) } +func (b *Bus) SpillRetention() time.Duration { return b.retention } +func (b *Bus) DSN() string { return b.cfg.DSN } +func (b *Bus) PurgeBefore(ctx context.Context, olderThan time.Duration) (int64, error) { + return b.purgeBefore(ctx, olderThan) +} diff --git a/core/services/pgbus/listen_ordering_internal_test.go b/core/services/pgbus/listen_ordering_internal_test.go new file mode 100644 index 000000000000..5e9dd0b6835f --- /dev/null +++ b/core/services/pgbus/listen_ordering_internal_test.go @@ -0,0 +1,176 @@ +package pgbus + +import ( + "context" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/testutil" +) + +// Internal on purpose. What these specs are about is the ORDER in which two +// registrations reach the connection, and that order is not observable from +// outside the package: a channel that is subscribed but not listened behaves +// exactly like a deployment where nobody is publishing. +// +// They drive that order through the barrier seam rather than by racing +// goroutines and hoping. The natural window here was measured at zero hits in +// forty attempts and ten out of ten once widened by hand, so a spec that races +// for it would pass by luck on a defect that leaves a whole subject root deaf +// until the connection drops. +var _ = Describe("ordering LISTEN and UNLISTEN on one root", func() { + var b *Bus + + newBus := func(cfg Config) *Bus { + GinkgoHelper() + built, err := New(context.Background(), cfg) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(built.Close) + return built + } + + BeforeEach(func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(Migrate(context.Background(), db)).To(Succeed()) + b = newBus(Config{DSN: dsn, DB: db}) + }) + + It("does not let a new LISTEN overtake the UNLISTEN it races", func() { + leaving, err := b.Subscribe("jobs.race.leaving", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + + unlistenDecided := make(chan struct{}) + release := make(chan struct{}) + arrived := make(chan struct{}) + listenIssued := make(chan struct{}, 1) + // Set after the first Subscribe, so only the racing pair is observed. + b.listenBarrier = func(stage, op string) { + switch { + case op == "UNLISTEN" && stage == "issue": + close(unlistenDecided) + <-release + case op == "LISTEN" && stage == "enter": + close(arrived) + case op == "LISTEN" && stage == "issue": + listenIssued <- struct{}{} + } + } + + unsubscribed := make(chan error, 1) + go func() { + defer GinkgoRecover() + unsubscribed <- leaving.Unsubscribe() + }() + // The last unsubscribe has decided to UNLISTEN and has not issued it. + // Eventually rather than a bare receive: a carrier that never reaches + // the decision at all must fail this spec, not hang it. + Eventually(unlistenDecided).Should(BeClosed()) + + delivered := make(chan []byte, 1) + subscribed := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := b.Subscribe("jobs.race.arriving", func(d []byte) { delivered <- d }) + subscribed <- err + }() + // The competing subscribe has started. Waiting for this rather than for + // a duration is what makes the assertion below decide something. + Eventually(arrived).Should(BeClosed()) + + // The bite, and it is deterministic in both directions. Unserialized, + // this LISTEN is on the wire within microseconds of the line above and + // the UNLISTEN then undoes it. Serialized, it provably cannot be + // issued while the unsubscribe holds the ordering lock. + Consistently(listenIssued).ShouldNot(Receive()) + + close(release) + Eventually(unsubscribed).Should(Receive(BeNil())) + Eventually(subscribed).Should(Receive(BeNil())) + + // The consequence, stated as delivery: the root must not be deaf. + Expect(b.Publish("jobs.race.arriving", map[string]string{"m": "after"})).To(Succeed()) + Eventually(delivered).Should(Receive(MatchJSON(`{"m":"after"}`))) + }) + + It("listens once for a root and stops listening when its last subscriber leaves", func() { + // The half of Unsubscribe that the delivery spec cannot see. Without + // the registration being removed from the map, the second unsubscribe + // never decides it was the last, and the connection keeps listening on + // a root nothing is subscribed to for the life of the process. + issued := make(chan string, 4) + b.listenBarrier = func(stage, op string) { + if stage == "issue" { + issued <- op + } + } + + one, err := b.Subscribe("jobs.refcount.a", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + two, err := b.Subscribe("jobs.refcount.b", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + + Expect(one.Unsubscribe()).To(Succeed()) + Expect(two.Unsubscribe()).To(Succeed()) + + Expect(issued).To(HaveLen(2)) + Expect(<-issued).To(Equal("LISTEN")) + Expect(<-issued).To(Equal("UNLISTEN")) + }) + + It("waits for the resolver as well as the listener when it closes", func() { + // Close returning while the resolver is still dispatching hands the + // caller a bus whose handlers can still fire after shutdown, against a + // database handle the process is about to drop. + // + // Asserted as "Close has not returned yet" rather than as "the done + // channel is closed afterwards": the resolver exits on its own when the + // context is cancelled, so the second shape passes whether Close waits + // or not. + dispatching := make(chan struct{}, 1) + release := make(chan struct{}) + var releaseOnce sync.Once + // Registered inside the It, so it runs BEFORE the bus teardown the + // BeforeEach registered. A failing assertion below must not leave the + // resolver parked on a channel that Close is waiting for. + DeferCleanup(func() { releaseOnce.Do(func() { close(release) }) }) + b.listenBarrier = func(stage, op string) { + if op == "DELIVER" && stage == "enter" { + select { + case dispatching <- struct{}{}: + default: + } + <-release + } + } + _, err := b.Subscribe("jobs.close.pending", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(b.Publish("jobs.close.pending", map[string]string{"m": "x"})).To(Succeed()) + Eventually(dispatching).Should(Receive()) + + closed := make(chan struct{}) + go func() { + defer GinkgoRecover() + b.Close() + close(closed) + }() + + Consistently(closed).ShouldNot(BeClosed()) + releaseOnce.Do(func() { close(release) }) + Eventually(closed).Should(BeClosed()) + }) + + It("stops the delivery goroutine of a subscription that has left", func() { + // The other half. A subscription that is deregistered but whose runner + // is never stopped leaks a goroutine per Unsubscribe, and the delivery + // spec cannot tell the two halves apart because either one alone makes + // the handler silent. + s, err := b.Subscribe("jobs.refcount.runner", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + + Expect(s.Unsubscribe()).To(Succeed()) + + Eventually(s.(*subscription).done).Should(BeClosed()) + }) +}) diff --git a/core/services/pgbus/listener.go b/core/services/pgbus/listener.go new file mode 100644 index 000000000000..f49a58924461 --- /dev/null +++ b/core/services/pgbus/listener.go @@ -0,0 +1,352 @@ +package pgbus + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + "github.com/mudler/xlog" +) + +// DefaultQueueDepth bounds how many notifications may wait for the resolver +// before this carrier drops them. +// +// It is a bound and not a blocking channel, on purpose. PostgreSQL keeps the +// notifications that were not delivered in one queue for the whole server, and +// it kills a session whose queue the client does not drain. A listener that +// blocked on a slow resolver would lose its connection, and with it every later +// broadcast. That is unbounded loss taken on to avoid bounded loss, and while it +// lasts it can block COMMIT for every other publisher on the server. +const DefaultQueueDepth = 1024 + +// DefaultSpillRetention is how long a spilled row outlives its notification +// before the purge loop deletes it. +// +// It is long compared with the delivery that it has to survive, which is one +// notification and one SELECT on the primary key, and short enough that a busy +// deployment does not collect the outputs of models in this table. +const DefaultSpillRetention = 10 * time.Minute + +// listenPollInterval bounds how long a LISTEN or UNLISTEN waits for the listener +// goroutine to leave WaitForNotification and run it. +// +// It bounds registration and never delivery: a notification wakes +// WaitForNotification at once. The connection survives the deadline, because pgx +// implements it as a read deadline and not as a close, so this costs one system +// call for each interval. +const listenPollInterval = 25 * time.Millisecond + +// subscribeTimeout bounds Subscribe when the listener is dialling a database +// that went away, so that a caller gets an error and not a goroutine that never +// returns. +const subscribeTimeout = 30 * time.Second + +// applicationNamePrefix marks the LISTEN sessions of this carrier in +// pg_stat_activity. An operator who counts the listeners, and a spec that must +// drop the connection of one carrier and not of its peer, need a session that +// they can name. +const applicationNamePrefix = "localai_pgbus_" + +// listenCmd is a LISTEN or UNLISTEN that is passed to the goroutine that owns the +// connection. The connection is not safe for concurrent use, and the listener +// goroutine waits in WaitForNotification most of the time, so registrations are +// queued to it and are not run against the connection. +type listenCmd struct { + sql string + done chan error +} + +// inbound is a notification as it came off the connection, before it is decoded, +// resolved and dispatched. +type inbound struct { + channel string + payload string +} + +// Dropped reports how many broadcasts this replica received and lost: in the +// queue between the listener and the resolver, in the queue of a handler, or +// because a notification could not be decoded or resolved. It only grows. The +// same events are counted, by stage, in localai_pgbus_dropped_total. +// +// It is for reports. Nothing may decide on it. A lost broadcast is a fact about a +// frontend, while the conditions that a scheduler acts on are all facts about a +// worker. A reaper that read it would report a failure of the carrier as the +// absence of a worker. +func (b *Bus) Dropped() uint64 { return b.dropped.Load() } + +// OnReconnect registers a callback that runs after the LISTEN connection was +// opened again and every channel was listened again. +// +// Consumers find it through an optional interface (interface{ OnReconnect(func()) }) +// and not through messaging.Broadcaster, so a carrier that cannot reconnect does +// not have to say that it can. Nothing fails to compile if the call is removed +// from here: a map that is kept in step through broadcasts would stop converging +// without any sign, so the spec in listener_test.go is the only guard. +// +// A nil callback is ignored. +func (b *Bus) OnReconnect(cb func()) { + if cb == nil { + return + } + b.cbMu.Lock() + b.reconnectCbs = append(b.reconnectCbs, cb) + b.cbMu.Unlock() +} + +// runReconnectCallbacks runs the registered callbacks, each on a goroutine of its +// own. +// +// They are detached on purpose. This runs on the goroutine that owns the LISTEN +// connection, and the callbacks read a whole durable source back from the +// database. To run one inline would put a query on the path whose only job is to +// keep the queue of PostgreSQL drained, and a callback that never returned would +// leave the carrier deaf for good. +// +// The slice is copied under the lock, so a callback that registers another does +// not deadlock. +func (b *Bus) runReconnectCallbacks() { + b.cbMu.Lock() + cbs := append([]func(){}, b.reconnectCbs...) + b.cbMu.Unlock() + for _, cb := range cbs { + go cb() + } +} + +// dial opens a LISTEN connection with the application name of this carrier. +// +// The DSN is parsed again for each dial. pgx may record the state of a +// connection in the config that it receives, and to use one config for several +// dials is not part of its contract. +func (b *Bus) dial(ctx context.Context) (*pgx.Conn, error) { + return dial(ctx, b.cfg.DSN, b.appName) +} + +func dial(ctx context.Context, dsn, appName string) (*pgx.Conn, error) { + connCfg, err := pgx.ParseConfig(dsn) + if err != nil { + return nil, fmt.Errorf("pgbus: parsing the LISTEN connection string: %w", err) + } + if connCfg.RuntimeParams == nil { + connCfg.RuntimeParams = map[string]string{} + } + connCfg.RuntimeParams["application_name"] = appName + return pgx.ConnectConfig(ctx, connCfg) +} + +// command passes a LISTEN or UNLISTEN to the goroutine that owns the connection +// and waits for it. To return before the server acknowledged it would lose every +// message published in the gap. +// +// A Subscribe that returned is therefore a registration that the server +// acknowledged. This holds when it issued nothing because the channel was +// already listened: because of listenMu, the Subscribe that did issue the LISTEN +// had its acknowledgement before this one could see its registration. +func (b *Bus) command(sql string) error { + cmd := listenCmd{sql: sql, done: make(chan error, 1)} + timeout := time.NewTimer(subscribeTimeout) + defer timeout.Stop() + + select { + case b.cmds <- cmd: + case <-b.ctx.Done(): + return errors.New("pgbus: the carrier is closed") + case <-timeout.C: + return fmt.Errorf("pgbus: %q timed out waiting for the listen connection", sql) + } + select { + case err := <-cmd.done: + return err + case <-b.ctx.Done(): + return errors.New("pgbus: the carrier is closed") + case <-timeout.C: + return fmt.Errorf("pgbus: %q timed out on the listen connection", sql) + } +} + +// listen owns the pinned connection for the life of the Bus. It is the only +// goroutine that uses it, which is what lets many callers share a connection that +// is not safe for concurrent use. +// +// A notification is copied off the connection here and is never processed here. +// Between the wait and offer, the loop must not do anything that can block on a +// database, on a handler or on the network, because whatever it waited for would +// hold the notifications of this session in a queue that it shares with every +// other publisher of the server. +func (b *Bus) listen(conn *pgx.Conn) { + defer close(b.listenerDone) + defer func() { + b.connected.Store(false) + // A fresh context, because b.ctx is already cancelled when this runs and + // the close still has to reach the server. + closeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = conn.Close(closeCtx) + }() + + for { + if b.ctx.Err() != nil { + return + } + // Registrations first: a subscriber that has waited must not be held up + // behind a poll interval that has just started. + select { + case cmd := <-b.cmds: + _, err := conn.Exec(b.ctx, cmd.sql) + cmd.done <- err + continue + default: + } + + waitCtx, cancel := context.WithTimeout(b.ctx, listenPollInterval) + n, err := conn.WaitForNotification(waitCtx) + deadline := waitCtx.Err() != nil + cancel() + + switch { + case b.ctx.Err() != nil: + return + case deadline: + // The usual case: nothing arrived in the poll window. pgx keeps the + // connection when its read deadline passes, so this is not a failure. + continue + case err != nil: + xlog.Warn("Broadcast carrier lost its listen connection, dialling again", "error", err) + replacement, ok := b.redial(conn) + if !ok { + return + } + conn = replacement + continue + } + b.offer(n.Channel, n.Payload) + } +} + +// offer hands a notification to the resolver. It never blocks: the only job of +// the listener is to keep the queue of PostgreSQL drained. +// +// The drop is counted as well as logged. The log is on the replica that took the +// loss, which is not always the party that needs to know, and a counter that only +// grows is what a spec and an operator can both read. +func (b *Bus) offer(channel, payload string) { + select { + case b.inbound <- inbound{channel: channel, payload: payload}: + default: + b.lose(stageListener) + xlog.Error("Broadcast carrier dropped a notification: the resolver is not keeping up", + "channel", channel, "depth", cap(b.inbound), "dropped", b.dropped.Load()) + } +} + +// resolve is where a spilled broadcast is read back and where every notification +// is dispatched. Both are off the listener on purpose, and both are on one +// goroutine, so a spilled message and an inline message of one subject keep the +// order in which they were published. +// +// This half may block. A notification is processed here and never on the +// connection that it arrived on, so a slow SELECT here costs this replica a +// bounded loss through offer and costs the server nothing. +func (b *Bus) resolve() { + defer close(b.resolverDone) + for { + select { + case <-b.ctx.Done(): + return + case in := <-b.inbound: + b.deliver(in.channel, in.payload) + } + } +} + +// redial replaces a dead LISTEN connection and restores every registration on it. +// Without the restore, the carrier comes back deaf, which looks like a deployment +// where nobody publishes any more. +func (b *Bus) redial(dead *pgx.Conn) (*pgx.Conn, bool) { + b.connected.Store(false) + closeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + _ = dead.Close(closeCtx) + cancel() + + backoff := 100 * time.Millisecond + for { + if b.ctx.Err() != nil { + return nil, false + } + conn, err := b.dial(b.ctx) + if err == nil { + if err = b.relisten(conn); err == nil { + b.connected.Store(true) + b.metrics.reconnect() + xlog.Info("Broadcast carrier reconnected") + // After the LISTEN and never before it. A callback reads from the + // durable source, and one that ran while the channels were not + // registered would converge against a carrier that cannot yet + // tell it about the next change. + b.runReconnectCallbacks() + return conn, true + } + closeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + _ = conn.Close(closeCtx) + cancel() + } + select { + case <-b.ctx.Done(): + return nil, false + case <-time.After(backoff): + } + if backoff < 5*time.Second { + backoff *= 2 + } + } +} + +// relisten registers again every channel that still has a subscriber. +// +// A carrier that reconnected without this is connected and deaf, and no error and +// no log line says so. +func (b *Bus) relisten(conn *pgx.Conn) error { + b.mu.Lock() + channels := make([]string, 0, len(b.subs)) + for channel, subs := range b.subs { + if len(subs) > 0 { + channels = append(channels, channel) + } + } + b.mu.Unlock() + + for _, channel := range channels { + if _, err := conn.Exec(b.ctx, "LISTEN "+pgx.Identifier{channel}.Sanitize()); err != nil { + return err + } + } + return nil +} + +// purge deletes the spilled rows that every replica has had time to read. +// +// Every replica purges. The DELETE is idempotent, and a deployment in which one +// replica purged would stop deleting when that replica went down. +func (b *Bus) purge() { + defer close(b.purgerDone) + interval := b.cfg.SweepInterval + if interval <= 0 { + // Half the retention, so a row is deleted within one retention after it + // ages out and not within two. + interval = b.retention / 2 + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-b.ctx.Done(): + return + case <-ticker.C: + if _, err := b.purgeBefore(b.ctx, b.retention); err != nil && b.ctx.Err() == nil { + xlog.Warn("Broadcast carrier could not retire spilled messages", "error", err) + } + } + } +} diff --git a/core/services/pgbus/listener_test.go b/core/services/pgbus/listener_test.go new file mode 100644 index 000000000000..a6ec4ba8eeec --- /dev/null +++ b/core/services/pgbus/listener_test.go @@ -0,0 +1,312 @@ +package pgbus_test + +import ( + "context" + "fmt" + "reflect" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// The two failure modes NATS did not have, and both are silent. +// +// Neither is provoked with a double or a clock here. A slow consumer is a real +// SELECT held by a real table lock, and a dropped listener is a real +// pg_terminate_backend against the session's own application_name, because a +// double that never touches the transport cannot fail the way the transport +// fails. +var _ = Describe("the listener under a slow consumer and a dropped connection", func() { + var ( + db *gorm.DB + dsn string + ) + + newBus := func(cfg pgbus.Config) *pgbus.Bus { + GinkgoHelper() + cfg.DSN, cfg.DB = dsn, db + b, err := pgbus.New(context.Background(), cfg) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(b.Close) + return b + } + + received := func(b *pgbus.Bus, filter string) <-chan []byte { + GinkgoHelper() + out := make(chan []byte, 128) + s, err := b.Subscribe(filter, func(data []byte) { out <- data }) + Expect(err).ToNot(HaveOccurred()) + Expect(s).ToNot(BeNil()) + return out + } + + // terminate drops exactly this carrier's LISTEN session and nobody else's. + // Matching on the application name rather than on "every backend except + // mine" is what lets a spec assert on a peer bus that stayed 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") + } + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + }) + + Describe("a resolver that cannot keep up", func() { + It("keeps draining the connection and drops, rather than stalling on the connection", func() { + // The hazard this whole split exists for. PostgreSQL holds + // undelivered notifications in an async queue that is shared by + // every session on the SERVER, so a listener that stopped draining + // would not merely fall behind: it can block COMMIT for unrelated + // publishers, and the backend eventually kills the session, which + // costs this replica every subsequent broadcast rather than these + // few. + // + // The resolver is stalled through the transport and not through a + // seam: a spilled broadcast is resolved with a SELECT on + // bus_messages, and an ACCESS EXCLUSIVE lock held by another + // transaction stops that SELECT for exactly as long as the spec + // wants it stopped. Nothing here waits on a clock. + const queue = 8 + const sent = queue + 64 + + b := newBus(pgbus.Config{Queue: queue, Retention: time.Hour}) + out := received(b, "jobs.spill.probe") + + channel, err := pgbus.ChannelFor("jobs.spill.probe") + Expect(err).ToNot(HaveOccurred()) + + // Distinct rows, so every notification carries a distinct payload: + // PostgreSQL collapses duplicate notifications raised in one + // transaction, and a spec that relied on identical payloads would + // be counting on that not happening. + ids := make([]string, sent) + for i := range ids { + ids[i] = fmt.Sprintf("spill-%02d", i) + Expect(db.Create(&pgbus.BusMessage{ + ID: ids[i], Subject: "jobs.spill.probe", Payload: []byte(fmt.Sprintf(`{"i":%d}`, i)), + }).Error).To(Succeed()) + } + + held := db.Begin() + Expect(held.Error).ToNot(HaveOccurred()) + released := false + release := func() { + if !released { + released = true + Expect(held.Rollback().Error).ToNot(HaveOccurred()) + } + } + DeferCleanup(release) + Expect(held.Exec("LOCK TABLE bus_messages IN ACCESS EXCLUSIVE MODE").Error).To(Succeed()) + + for _, id := range ids { + Expect(db.Exec("SELECT pg_notify(?, ?)", channel, + fmt.Sprintf(`{"s":"jobs.spill.probe","i":%q}`, id)).Error).To(Succeed()) + } + + // The listener took every notification off the connection while the + // resolver could not move: the queue filled and the surplus was + // dropped. A listener that had blocked instead would leave this at + // zero forever. + Eventually(b.Dropped, 30*time.Second).Should(BeNumerically(">", 0)) + Expect(b.IsConnected()).To(BeTrue(), "dropping must not cost the connection") + + release() + + // Still a working carrier once the consumer catches up. Bounded + // loss, not a wedge. + Eventually(out, 30*time.Second).Should(Receive()) + }) + + It("reports nothing dropped on a carrier that is keeping up", func() { + // The zero the spec above is measured against. Without it a Dropped + // that returned a constant would satisfy the assertion there. + b := newBus(pgbus.Config{}) + out := received(b, "jobs.fast.consumer") + + Expect(b.Publish("jobs.fast.consumer", map[string]int{"i": 1})).To(Succeed()) + + Eventually(out).Should(Receive()) + Expect(b.Dropped()).To(BeZero()) + }) + }) + + Describe("losing the LISTEN connection", func() { + It("invokes every OnReconnect callback once the channels are listened again", func() { + // syncstate.SyncedMap syncs deltas only, so this callback is the one + // mechanism it has to converge after a gap. It reaches it through an + // optional interface assertion, which means deleting the invocation + // compiles, passes every other spec, and leaves every adopter + // silently diverging. This spec is the only guard there is. + b := newBus(pgbus.Config{}) + _ = received(b, "jobs.reconnect.probe") + + calls := make(chan struct{}, 8) + b.OnReconnect(func() { calls <- struct{}{} }) + // Not fired by the connection it was registered on: a callback that + // ran on the first connect would re-hydrate every adopter at boot + // for no reason. + Consistently(calls, 200*time.Millisecond).ShouldNot(Receive()) + + terminate(b) + + // The callback cannot fire without a lost connection, a redial and a + // completed re-LISTEN, so its arrival is what proves the whole cycle + // ran. Sampling IsConnected for the false half would be sampling: the + // window between the loss and the first redial is sub-millisecond, + // and a spec that catches it nine times in ten is worse than none. + Eventually(calls, 30*time.Second).Should(Receive()) + Expect(b.IsConnected()).To(BeTrue()) + }) + + It("delivers again on a channel it had to re-listen", func() { + // Separate from the callback assertion on purpose: a carrier can be + // reconnected, and report itself connected, and be deaf, because a + // connection without its LISTEN registrations produces no error + // anywhere. Deleting the re-LISTEN reddens this and leaves the + // callback spec above green. + b := newBus(pgbus.Config{}) + publisher := newBus(pgbus.Config{}) + out := received(b, "jobs.relisten.probe") + + terminate(b) + + Eventually(func() bool { + // Retried because a NOTIFY that lands between the loss and the + // re-LISTEN is genuinely gone: this carrier is at-most-once, and + // what is being asserted is that delivery RESUMES. + _ = publisher.Publish("jobs.relisten.probe", map[string]string{"m": "after"}) + select { + case <-out: + return true + default: + return false + } + }, 30*time.Second, 200*time.Millisecond).Should(BeTrue()) + }) + + It("does not run a reconnect callback on the goroutine that owns the connection", func() { + // A callback re-hydrates from a durable source, which is a database + // query. Running one inline would put that query on the path whose + // only job is to keep PostgreSQL's async queue draining, and a + // callback that never returned would leave this carrier deaf for + // good. Deleting the `go` compiles and is invisible everywhere else. + b := newBus(pgbus.Config{}) + publisher := newBus(pgbus.Config{}) + out := received(b, "jobs.blockedcb.probe") + + entered := make(chan struct{}, 1) + forever := make(chan struct{}) + DeferCleanup(func() { close(forever) }) + b.OnReconnect(func() { + entered <- struct{}{} + <-forever + }) + + terminate(b) + + Eventually(entered, 30*time.Second).Should(Receive()) + Eventually(func() bool { + _ = publisher.Publish("jobs.blockedcb.probe", map[string]string{"m": "after"}) + select { + case <-out: + return true + default: + return false + } + }, 30*time.Second, 200*time.Millisecond).Should(BeTrue()) + }) + }) + + Describe("naming the LISTEN session", func() { + It("gives each carrier an application_name an operator can find in pg_stat_activity", func() { + // Without it, dropping one replica's listener means dropping every + // backend on the database, and an operator asking how many replicas + // are listening has nothing to count. + b := newBus(pgbus.Config{}) + + var sessions int64 + Expect(db.Raw( + "SELECT count(*) FROM pg_stat_activity WHERE application_name = ?", b.ApplicationName(), + ).Scan(&sessions).Error).To(Succeed()) + + Expect(b.ApplicationName()).To(HavePrefix("localai_pgbus_")) + Expect(sessions).To(BeNumerically("==", 1)) + }) + + It("names two carriers on one database differently", func() { + first := newBus(pgbus.Config{}) + second := newBus(pgbus.Config{}) + + Expect(first.ApplicationName()).ToNot(Equal(second.ApplicationName())) + }) + }) + + Describe("the queue and retention settings", func() { + It("takes the exported defaults when both are zero", func() { + b := newBus(pgbus.Config{}) + + Expect(b.QueueDepth()).To(Equal(pgbus.DefaultQueueDepth)) + Expect(b.SpillRetention()).To(Equal(pgbus.DefaultSpillRetention)) + }) + + It("takes what it was configured with when they are not", func() { + b := newBus(pgbus.Config{Queue: 7, Retention: 3 * time.Minute}) + + Expect(b.QueueDepth()).To(Equal(7)) + Expect(b.SpillRetention()).To(Equal(3 * time.Minute)) + }) + + It("runs the purge loop off the retention when no interval is configured", func() { + // The interval is derived rather than configured in production, so + // nothing else proves the derivation is wired: a purge loop pinned + // to a constant would retire rows too, just not on the schedule the + // retention asks for. + Expect(db.Create(&pgbus.BusMessage{ + ID: "aged", Subject: "jobs.x", Payload: []byte(`{}`), + }).Error).To(Succeed()) + Expect(db.Exec( + "UPDATE bus_messages SET created_at = now() - interval '1 hour' WHERE id = 'aged'", + ).Error).To(Succeed()) + + newBus(pgbus.Config{Retention: 40 * time.Millisecond}) + + Eventually(func() (int64, error) { + var n int64 + err := db.Model(&pgbus.BusMessage{}).Count(&n).Error + return n, err + }, 30*time.Second).Should(BeZero()) + }) + }) + + Describe("what the carrier's own health is allowed to reach", func() { + It("keeps Dropped and IsConnected off the interface consumers hold", func() { + // The invariant, guarded structurally. "This replica is behind" and + // "this replica's listener is down" are facts about a FRONTEND; the + // conditions a scheduler acts on are all facts about a WORKER. As + // long as neither method is reachable through messaging.Broadcaster, + // no consumer can read a carrier failure as a worker's absence + // without reaching for the concrete type and saying so. + carrier := reflect.TypeOf((*messaging.Broadcaster)(nil)).Elem() + + _, hasDropped := carrier.MethodByName("Dropped") + _, hasConnected := carrier.MethodByName("IsConnected") + + Expect(hasDropped).To(BeFalse(), "a subscriber's loss must not be readable as a node's state") + Expect(hasConnected).To(BeFalse(), "a carrier's link must not be readable as a node's state") + }) + }) +}) diff --git a/core/services/pgbus/metrics.go b/core/services/pgbus/metrics.go new file mode 100644 index 000000000000..08faa628ce72 --- /dev/null +++ b/core/services/pgbus/metrics.go @@ -0,0 +1,84 @@ +package pgbus + +import ( + "context" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// Where a broadcast was lost. The value is the "stage" attribute of the +// dropped counter. +const ( + // stageListener: the notification queue between the LISTEN connection and + // the resolver was full. + stageListener = "listener" + // stageSubscription: the queue of one handler was full. + stageSubscription = "subscription" + // stageResolve: the notification could not be decoded, or the row that it + // names could not be read. + stageResolve = "resolve" +) + +// Which way a broadcast travelled. The value is the "path" attribute of the +// published counter. +const ( + pathInline = "inline" + pathSpill = "spill" +) + +// metrics are the counters of one Bus. +// +// A broadcast is at-most-once, so the counters are the only place where a loss +// is visible to an operator. Without them a loss shows up as a log line on the +// replica that took it, and as a missing event on the others. +type metrics struct { + published metric.Int64Counter + dropped metric.Int64Counter + spillPurged metric.Int64Counter + reconnections metric.Int64Counter +} + +// newMetrics creates the instruments on meter. An instrument that cannot be +// created is left nil and its updates are skipped, because a metric must not +// stop the carrier. +func newMetrics(meter metric.Meter) *metrics { + if meter == nil { + meter = otel.Meter("github.com/mudler/LocalAI") + } + m := &metrics{} + m.published, _ = meter.Int64Counter("localai_pgbus_published_total", + metric.WithDescription("Broadcasts published through the database, by path: inline in the notification, or spilled to a row")) + m.dropped, _ = meter.Int64Counter("localai_pgbus_dropped_total", + metric.WithDescription("Broadcasts that this replica received and lost, by stage: listener queue, subscription queue or resolve")) + m.spillPurged, _ = meter.Int64Counter("localai_pgbus_spill_purged_total", + metric.WithDescription("Spilled broadcast rows deleted after the retention")) + m.reconnections, _ = meter.Int64Counter("localai_pgbus_reconnections_total", + metric.WithDescription("Times the LISTEN connection was lost and opened again")) + return m +} + +func (m *metrics) publish(path string) { + if m.published != nil { + m.published.Add(context.Background(), 1, metric.WithAttributes(attribute.String("path", path))) + } +} + +func (m *metrics) drop(stage string) { + if m.dropped != nil { + m.dropped.Add(context.Background(), 1, metric.WithAttributes(attribute.String("stage", stage))) + } +} + +func (m *metrics) purge(rows int64) { + if m.spillPurged != nil && rows > 0 { + m.spillPurged.Add(context.Background(), rows) + } +} + +func (m *metrics) reconnect() { + if m.reconnections != nil { + m.reconnections.Add(context.Background(), 1) + } +} diff --git a/core/services/pgbus/metrics_test.go b/core/services/pgbus/metrics_test.go new file mode 100644 index 000000000000..36173b055c78 --- /dev/null +++ b/core/services/pgbus/metrics_test.go @@ -0,0 +1,181 @@ +package pgbus_test + +import ( + "context" + "fmt" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// A broadcast is at-most-once, so the counters are the only place where a loss is +// visible to an operator. Each spec makes the loss happen on a real database and +// reads the counter back through the OpenTelemetry SDK. +var _ = Describe("the counters of the PostgreSQL broadcast carrier", func() { + var ( + db *gorm.DB + dsn string + reader *sdkmetric.ManualReader + ) + + newBus := func(cfg pgbus.Config) *pgbus.Bus { + GinkgoHelper() + cfg.DSN, cfg.DB = dsn, db + cfg.Meter = sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)).Meter("pgbus-spec") + b, err := pgbus.New(context.Background(), cfg) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(b.Close) + return b + } + + // counter returns the sum of the counter with the given name over the data + // points that carry the attribute. An empty key sums them all. + counter := func(name, key, value string) int64 { + GinkgoHelper() + var rm metricdata.ResourceMetrics + Expect(reader.Collect(context.Background(), &rm)).To(Succeed()) + var sum int64 + for _, scope := range rm.ScopeMetrics { + for _, m := range scope.Metrics { + data, ok := m.Data.(metricdata.Sum[int64]) + if m.Name != name || !ok { + continue + } + for _, p := range data.DataPoints { + if key == "" { + sum += p.Value + continue + } + if v, found := p.Attributes.Value(attribute.Key(key)); found && v.AsString() == value { + sum += p.Value + } + } + } + } + return sum + } + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + reader = sdkmetric.NewManualReader() + }) + + It("counts a broadcast by the path that it took", func() { + b := newBus(pgbus.Config{}) + + Expect(b.Publish("jobs.m1.progress", map[string]string{"k": "small"})).To(Succeed()) + Expect(b.Publish("jobs.m1.progress", map[string]string{"k": "small"})).To(Succeed()) + Expect(b.Publish("jobs.m1.result", map[string]string{"k": strings.Repeat("a", 64*1024)})).To(Succeed()) + + Expect(counter("localai_pgbus_published_total", "path", "inline")).To(Equal(int64(2))) + Expect(counter("localai_pgbus_published_total", "path", "spill")).To(Equal(int64(1))) + }) + + It("does not count a broadcast that was refused", func() { + b := newBus(pgbus.Config{}) + + Expect(b.Publish("nodes.n1.backend.stop", "x")).ToNot(Succeed()) + + Expect(counter("localai_pgbus_published_total", "", "")).To(BeZero()) + }) + + It("counts a notification that the listener had to drop, by stage", func() { + b := newBus(pgbus.Config{Queue: 4, Retention: time.Hour}) + out := make(chan []byte, 128) + _, err := b.Subscribe("jobs.m2.probe", func(data []byte) { out <- data }) + Expect(err).ToNot(HaveOccurred()) + + const sent = 40 + for i := 0; i < sent; i++ { + Expect(db.Create(&pgbus.BusMessage{ + ID: fmt.Sprintf("m2-%02d", i), Subject: "jobs.m2.probe", Payload: []byte(`{}`), + }).Error).To(Succeed()) + } + // A lock that stops the SELECT of the resolver, so that the queue fills. + held := db.Begin() + Expect(held.Error).ToNot(HaveOccurred()) + released := false + release := func() { + if !released { + released = true + Expect(held.Rollback().Error).ToNot(HaveOccurred()) + } + } + DeferCleanup(release) + Expect(held.Exec("LOCK TABLE bus_messages IN ACCESS EXCLUSIVE MODE").Error).To(Succeed()) + for i := 0; i < sent; i++ { + Expect(db.Exec("SELECT pg_notify('localai_jobs', ?)", + fmt.Sprintf(`{"s":"jobs.m2.probe","i":"m2-%02d"}`, i)).Error).To(Succeed()) + } + + Eventually(func() int64 { return counter("localai_pgbus_dropped_total", "stage", "listener") }, + 30*time.Second).Should(BeNumerically(">", 0)) + Expect(int64(b.Dropped())).To(BeNumerically(">=", counter("localai_pgbus_dropped_total", "stage", "listener"))) + release() + }) + + It("counts a notification that could not be resolved, by stage", func() { + b := newBus(pgbus.Config{}) + out := make(chan []byte, 4) + _, err := b.Subscribe("jobs.m3.probe", func(data []byte) { out <- data }) + Expect(err).ToNot(HaveOccurred()) + + Expect(db.Exec("SELECT pg_notify('localai_jobs', ?)", `{"s":"jobs.m3.probe","i":"no-such-row"}`).Error).To(Succeed()) + Expect(db.Exec("SELECT pg_notify('localai_jobs', ?)", `not json`).Error).To(Succeed()) + + Eventually(func() int64 { return counter("localai_pgbus_dropped_total", "stage", "resolve") }, + 10*time.Second).Should(Equal(int64(2))) + Expect(out).ToNot(Receive()) + }) + + It("counts a message that a slow subscriber lost, by stage", func() { + b := newBus(pgbus.Config{}) + block := make(chan struct{}) + DeferCleanup(func() { close(block) }) + _, err := b.Subscribe("jobs.m4.slow", func([]byte) { <-block }) + Expect(err).ToNot(HaveOccurred()) + + for i := 0; i < 400; i++ { + Expect(b.Publish("jobs.m4.slow", map[string]int{"i": i})).To(Succeed()) + } + + Eventually(func() int64 { return counter("localai_pgbus_dropped_total", "stage", "subscription") }, + 30*time.Second).Should(BeNumerically(">", 0)) + }) + + It("counts the spilled rows that the purge deleted", func() { + b := newBus(pgbus.Config{Retention: time.Hour}) + for _, id := range []string{"p1", "p2", "p3"} { + Expect(db.Create(&pgbus.BusMessage{ID: id, Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed()) + } + Expect(db.Exec("UPDATE bus_messages SET created_at = now() - interval '2 hours'").Error).To(Succeed()) + + deleted, err := b.PurgeBefore(context.Background(), time.Hour) + + Expect(err).ToNot(HaveOccurred()) + Expect(deleted).To(Equal(int64(3))) + Expect(counter("localai_pgbus_spill_purged_total", "", "")).To(Equal(int64(3))) + }) + + It("counts a reconnection of the LISTEN connection", func() { + b := newBus(pgbus.Config{}) + _, err := b.Subscribe("jobs.m5.probe", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(counter("localai_pgbus_reconnections_total", "", "")).To(BeZero()) + + Expect(db.Exec("SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE application_name = ?", b.ApplicationName()).Error).To(Succeed()) + + Eventually(func() int64 { return counter("localai_pgbus_reconnections_total", "", "") }, + 30*time.Second, 100*time.Millisecond).Should(Equal(int64(1))) + }) +}) diff --git a/core/services/pgbus/pgbus_suite_test.go b/core/services/pgbus/pgbus_suite_test.go new file mode 100644 index 000000000000..42973abb027b --- /dev/null +++ b/core/services/pgbus/pgbus_suite_test.go @@ -0,0 +1,13 @@ +package pgbus_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestPgbus(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "pgbus Suite") +} diff --git a/core/services/pgbus/spill.go b/core/services/pgbus/spill.go new file mode 100644 index 000000000000..69053e6fff0f --- /dev/null +++ b/core/services/pgbus/spill.go @@ -0,0 +1,49 @@ +package pgbus + +import ( + "context" + "fmt" + "time" + + "gorm.io/gorm" +) + +// spillSweepSQL deletes the spilled broadcasts that are older than the +// retention it is given. +// +// The database computes the cutoff and Go does not. Replicas do not share a +// clock: a replica that runs minutes ahead would compute a cutoff in the future +// and delete rows that its peers have not read yet, and a replica that runs +// behind would never delete anything. now() is the one clock that every replica +// agrees on, because it is the clock that stamped the rows. +// +// It is PostgreSQL only, like the rest of this package. New refuses a handle on +// any other dialect. +const spillSweepSQL = "DELETE FROM bus_messages WHERE created_at < now() - make_interval(secs => ?)" + +// BusMessage is one broadcast that is too large for a notification. +type BusMessage struct { + ID string `gorm:"primaryKey;size:36"` + Subject string `gorm:"size:255;index"` + Payload []byte + CreatedAt time.Time `gorm:"index"` +} + +// Migrate creates the spill table. It is separate from New because a deployment +// migrates once while every replica opens a carrier. +func Migrate(ctx context.Context, db *gorm.DB) error { + return db.WithContext(ctx).AutoMigrate(&BusMessage{}) +} + +// purgeBefore deletes the spilled rows older than olderThan and returns how many +// it deleted. The count is the point: a purge loop that retires nothing reports +// success as well, and a spill table that grows without bound fills the disk +// long after the change that caused it. +func (b *Bus) purgeBefore(ctx context.Context, olderThan time.Duration) (int64, error) { + res := b.cfg.DB.WithContext(ctx).Exec(spillSweepSQL, olderThan.Seconds()) + if res.Error != nil { + return 0, fmt.Errorf("pgbus: retiring spilled broadcasts: %w", res.Error) + } + b.metrics.purge(res.RowsAffected) + return res.RowsAffected, nil +} diff --git a/core/services/pgbus/spill_test.go b/core/services/pgbus/spill_test.go new file mode 100644 index 000000000000..9b89859f0cb3 --- /dev/null +++ b/core/services/pgbus/spill_test.go @@ -0,0 +1,272 @@ +package pgbus_test + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// PostgreSQL refuses a pg_notify payload of 8000 bytes or more: async.c checks +// strlen(payload) >= NOTIFY_PAYLOAD_MAX_LENGTH, so 7999 is the largest that is +// accepted. Verified against postgres:16 rather than read off the docs, because +// the docs say "less than 8000 bytes" and the off-by-one is the whole subject of +// these specs. +// +// This literal is deliberately NOT pgbus.MaxNotifyPayloadBytes. A size spec +// that measures itself against the constant it is testing passes for every +// value of that constant, which makes it a spec about nothing. +const notifyCap = 8000 + +// encodedSize is what the carrier puts on the wire for an inline broadcast: +// the subject, the JSON escaping and the envelope's own keys all travel with +// the caller's payload. A payload sized against the cap alone would ship 7999 +// bytes of data in an 8035-byte notification and fail at the server. +func encodedSize(subject string, data any) int { + GinkgoHelper() + subjectJSON, err := json.Marshal(subject) + Expect(err).ToNot(HaveOccurred()) + dataJSON, err := json.Marshal(data) + Expect(err).ToNot(HaveOccurred()) + return len(fmt.Sprintf(`{"s":%s,"d":%s}`, subjectJSON, dataJSON)) +} + +// payloadEncodingTo builds a valid single-key JSON object whose encoded +// notification is exactly total bytes. Valid on purpose: an oversized payload +// that the decoder would reject anyway cannot prove the SIZE routed it. +func payloadEncodingTo(subject string, total int) map[string]string { + GinkgoHelper() + data := map[string]string{"k": ""} + filler := total - encodedSize(subject, data) + Expect(filler).To(BeNumerically(">", 0)) + data["k"] = strings.Repeat("a", filler) + Expect(encodedSize(subject, data)).To(Equal(total)) + return data +} + +var _ = Describe("broadcasts too large for a notification", func() { + const subject = "jobs.size.probe" + + var ( + db *gorm.DB + dsn string + pub, sub *pgbus.Bus + ) + + spilledRows := func() int64 { + GinkgoHelper() + var n int64 + Expect(db.Model(&pgbus.BusMessage{}).Where("subject = ?", subject).Count(&n).Error).To(Succeed()) + return n + } + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + for _, b := range []**pgbus.Bus{&pub, &sub} { + built, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(built.Close) + *b = built + } + }) + + deliver := func(data any) []byte { + GinkgoHelper() + out := make(chan []byte, 4) + _, err := sub.Subscribe(subject, func(b []byte) { out <- b }) + Expect(err).ToNot(HaveOccurred()) + + Expect(pub.Publish(subject, data)).To(Succeed()) + + var got []byte + Eventually(out, 10*time.Second).Should(Receive(&got)) + return got + } + + // The three rows of the size table. Each states an ABSOLUTE size, so a + // mutation of the carrier's constant moves the behaviour without moving + // the expectation. + It("keeps a broadcast one byte under the cap in the notification itself", func() { + data := payloadEncodingTo(subject, notifyCap-1) + + got := deliver(data) + + Expect(got).To(MatchJSON(mustJSON(data))) + Expect(spilledRows()).To(BeZero(), "a notification PostgreSQL accepts must not touch the spill table") + }) + + It("spills a broadcast exactly at the cap, which PostgreSQL already refuses", func() { + data := payloadEncodingTo(subject, notifyCap) + + got := deliver(data) + + Expect(got).To(MatchJSON(mustJSON(data))) + Expect(spilledRows()).To(Equal(int64(1))) + }) + + It("carries a megabyte through the spill table byte for byte", func() { + // Absolute, not derived from the cap: this row must stay green under a + // changed constant, which separates "the spill path works" from "the spec + // sized itself against the thing it is testing". + data := map[string]string{"k": strings.Repeat("a", 1<<20)} + + got := deliver(data) + + Expect(got).To(Equal(mustJSON(data))) + Expect(spilledRows()).To(Equal(int64(1))) + }) + + It("stores the caller's payload in the spill row, not the envelope", func() { + data := payloadEncodingTo(subject, notifyCap) + + _ = deliver(data) + + var row pgbus.BusMessage + Expect(db.Where("subject = ?", subject).First(&row).Error).To(Succeed()) + Expect(row.Payload).To(Equal(mustJSON(data))) + Expect(row.ID).ToNot(BeEmpty()) + Expect(row.CreatedAt).ToNot(BeZero()) + }) + + It("delivers nothing for a notification that carries no payload at all", func() { + // Publish cannot produce this, because json.Marshal is never empty. It + // is reachable if anything else ever notifies on a localai_ channel, + // and a handler called with nil is a message that says nothing rather + // than no message, which is exactly the confusion this carrier must + // never create. The sentinel proves it kept carrying. + out := make(chan []byte, 4) + _, err := sub.Subscribe(subject, func(b []byte) { out <- b }) + Expect(err).ToNot(HaveOccurred()) + + Expect(db.Exec("SELECT pg_notify(?, ?)", "localai_jobs", + fmt.Sprintf(`{"s":%q}`, subject)).Error).To(Succeed()) + Expect(pub.Publish(subject, map[string]string{"m": "sentinel"})).To(Succeed()) + + var got []byte + Eventually(out, 10*time.Second).Should(Receive(&got)) + Expect(got).To(MatchJSON(`{"m":"sentinel"}`)) + }) + + It("delivers nothing, and keeps carrying, when a spilled row cannot be found", func() { + // A notification whose row is gone is a lost message, not an empty one: + // handing a handler nil would let a consumer read a carrier failure as + // a fact about the deployment. The sentinel proves the carrier is still + // alive afterwards, which needs no clock. + out := make(chan []byte, 4) + _, err := sub.Subscribe(subject, func(b []byte) { out <- b }) + Expect(err).ToNot(HaveOccurred()) + + Expect(db.Exec("SELECT pg_notify(?, ?)", "localai_jobs", + fmt.Sprintf(`{"s":%q,"i":"no-such-row"}`, subject)).Error).To(Succeed()) + Expect(pub.Publish(subject, map[string]string{"m": "sentinel"})).To(Succeed()) + + var got []byte + Eventually(out, 10*time.Second).Should(Receive(&got)) + Expect(got).To(MatchJSON(`{"m":"sentinel"}`)) + }) +}) + +var _ = Describe("retiring spilled rows", func() { + var ( + db *gorm.DB + dsn string + ) + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + }) + + // aged writes one spilled row and backdates it on the DATABASE clock, which + // is the clock the sweep compares against. + aged := func(id string) { + GinkgoHelper() + Expect(db.Create(&pgbus.BusMessage{ID: id, Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed()) + Expect(db.Exec("UPDATE bus_messages SET created_at = now() - interval '1 hour' WHERE id = ?", id).Error).To(Succeed()) + } + + rows := func() []string { + GinkgoHelper() + var ids []string + Expect(db.Model(&pgbus.BusMessage{}).Pluck("id", &ids).Error).To(Succeed()) + return ids + } + + It("leaves the cutoff to the database clock", func() { + // Pinned as a statement shape rather than as behaviour on purpose. The + // test container shares this host's clock, so a cutoff computed in Go + // and one computed by the server agree to the microsecond and no + // behavioural spec can tell them apart. The deployment they differ in + // is the one that matters: replicas whose clocks are minutes apart, + // where a Go-side cutoff reaps rows other replicas have not read yet. + Expect(pgbus.SpillSweepSQL).To(ContainSubstring("now()")) + Expect(pgbus.SpillSweepSQL).ToNot(ContainSubstring("created_at < ?")) + }) + + It("runs the sweep on its own, without anyone asking it to", func() { + // The statement and PurgeBefore are tested directly, and neither proves + // that the carrier ever calls them. A deleted `go` statement would leave + // the suite green and bus_messages growing for ever. The interval is a + // Config field so that this spec can exist. + aged("swept") + Expect(db.Create(&pgbus.BusMessage{ID: "kept", Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed()) + + b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, SweepInterval: 20 * time.Millisecond}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(b.Close) + + Eventually(rows, 30*time.Second).Should(ConsistOf("kept")) + }) + + Describe("PurgeBefore", func() { + var b *pgbus.Bus + + BeforeEach(func() { + // A retention long enough that the carrier's own purge loop cannot + // fire inside a spec, so what these assert on is the call they make + // and not a tick that happened to land. + var err error + b, err = pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, Retention: time.Hour}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(b.Close) + }) + + It("deletes a spilled row and says how many it deleted", func() { + // The count is the point: a purge loop that retires nothing and one + // that retires everything both look like a statement that ran. + aged("gone") + + deleted, err := b.PurgeBefore(context.Background(), 0) + + Expect(err).ToNot(HaveOccurred()) + Expect(deleted).To(BeNumerically("==", 1)) + Expect(rows()).To(BeEmpty()) + }) + + It("leaves a row younger than the cutoff alone", func() { + Expect(db.Create(&pgbus.BusMessage{ID: "young", Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed()) + + deleted, err := b.PurgeBefore(context.Background(), time.Hour) + + Expect(err).ToNot(HaveOccurred()) + Expect(deleted).To(BeZero()) + Expect(rows()).To(ConsistOf("young")) + }) + }) +}) + +func mustJSON(v any) []byte { + GinkgoHelper() + b, err := json.Marshal(v) + Expect(err).ToNot(HaveOccurred()) + return b +} diff --git a/go.mod b/go.mod index e16d680f5247..342faf887ce0 100644 --- a/go.mod +++ b/go.mod @@ -27,6 +27,7 @@ require ( github.com/gpustack/gguf-parser-go v0.25.0 github.com/hpcloud/tail v1.0.0 github.com/ipfs/go-log v1.0.5 + github.com/jackc/pgx/v5 v5.9.2 github.com/jaypipes/ghw v0.24.0 github.com/joho/godotenv v1.5.1 github.com/klauspost/cpuid/v2 v2.3.0 @@ -249,7 +250,6 @@ require ( github.com/google/go-querystring v1.1.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect - github.com/jackc/pgx/v5 v5.9.2 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect github.com/jung-kurt/gofpdf v1.16.2 // indirect From 029e5a0becbff470ab39feffff2a4a4e8f442798 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:32:07 +0000 Subject: [PATCH 015/162] fix(pgbus): read spilled broadcasts with several readers and keep the order A broadcast of 64 KiB is written to a row, and every replica reads that row with a SELECT. One reader at a time limited a replica to about a thousand such broadcasts per second, and a burst above that was dropped. Read up to eight rows at the same time and hand the notifications to the handlers in the order they arrived, so the order of one subject does not change. In a burst of 4000 broadcasts of 64 KiB from 16 publishers, the replica with one reader delivered 1860, and the replica with eight readers delivered all of them. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/pgbus/bus.go | 93 +++++++++++------ core/services/pgbus/listener.go | 66 +++++++++++- core/services/pgbus/pipeline_test.go | 149 +++++++++++++++++++++++++++ 3 files changed, 273 insertions(+), 35 deletions(-) create mode 100644 core/services/pgbus/pipeline_test.go diff --git a/core/services/pgbus/bus.go b/core/services/pgbus/bus.go index 2a20ec147ad8..e5b7937650f4 100644 --- a/core/services/pgbus/bus.go +++ b/core/services/pgbus/bus.go @@ -18,9 +18,12 @@ // a hint to go and look. // // Limits, measured on a four-core database: about 16,000 broadcasts a second of -// 200 bytes, about 8,000 a second with 20 listeners or with 2 to 8 KiB, and -// about 1,000 a second of 64 KiB. A broadcast that does not fit in a -// notification is written to a row and costs a SELECT on every replica. +// 200 bytes, and about 8,000 a second with 20 listeners or with 2 to 8 KiB. A +// broadcast that does not fit in a notification is written to a row and costs a +// SELECT on every replica; a replica reads several rows at the same time. With +// one reader at a time, a replica lost half of a burst of 64 KiB broadcasts +// above about 1,000 a second. Above the limits, broadcasts are dropped and +// counted, and never queued without a bound. package pgbus import ( @@ -434,54 +437,84 @@ func (b *Bus) barrier(stage, op string) { } } -// deliver resolves one notification and hands it to the subscribers whose -// filters match. -func (b *Bus) deliver(channel, payload string) { - b.barrier("enter", "DELIVER") +// pending is one notification on its way from the listener to the handlers. The +// resolver decodes it in the order it arrived, and a spilled one is read back by +// a worker while the next ones are decoded. The dispatcher takes the pending +// notifications in the same order and waits for each one to be ready. +type pending struct { + channel string + n notification + data []byte + // ok is false when the notification must not be delivered. + ok bool + ready chan struct{} +} - var n notification - if err := json.Unmarshal([]byte(payload), &n); err != nil { +// decode turns a notification as it came off the connection into a pending one. +// A spilled one is not ready: the caller starts fetch for it. +func (b *Bus) decode(in inbound) *pending { + p := &pending{channel: in.channel, ready: make(chan struct{})} + if err := json.Unmarshal([]byte(in.payload), &p.n); err != nil { b.lose(stageResolve) - xlog.Error("Broadcast carrier received an undecodable notification", "channel", channel, "error", err) - return + xlog.Error("Broadcast carrier received an undecodable notification", "channel", in.channel, "error", err) + close(p.ready) + return p } - - data := []byte(n.Data) - if n.SpillID == "" && len(data) == 0 { + p.data = []byte(p.n.Data) + if p.n.SpillID == "" && len(p.data) == 0 { // Publish cannot produce this, because json.Marshal never returns an // empty encoding. It is possible only if something other than this // carrier notifies on one of its channels. A handler that received nil // would see a message that says nothing, and no message is better. b.lose(stageResolve) - xlog.Error("Broadcast carrier received a notification with no payload", "channel", channel, "subject", n.Subject) + xlog.Error("Broadcast carrier received a notification with no payload", "channel", in.channel, "subject", p.n.Subject) + close(p.ready) + return p + } + p.ok = true + if p.n.SpillID == "" { + close(p.ready) + } + return p +} + +// fetch reads the row of a spilled notification and marks it ready. +func (b *Bus) fetch(p *pending) { + defer close(p.ready) + resolved, err := b.resolveSpill(p.n.SpillID) + if err != nil { + // A dropped broadcast and nothing more. A handler that got an empty + // message instead would read a carrier failure as a fact about the + // deployment. + p.ok = false + b.lose(stageResolve) + xlog.Error("Broadcast carrier could not resolve a spilled message", "subject", p.n.Subject, "id", p.n.SpillID, "error", err) return } - if n.SpillID != "" { - resolved, err := b.resolveSpill(n.SpillID) - if err != nil { - // A dropped broadcast and nothing more. A handler that got an empty - // message instead would read a carrier failure as a fact about the - // deployment. - b.lose(stageResolve) - xlog.Error("Broadcast carrier could not resolve a spilled message", "subject", n.Subject, "id", n.SpillID, "error", err) - return - } - data = resolved + p.data = resolved +} + +// deliver hands a notification that is ready to the subscribers whose filters +// match. +func (b *Bus) deliver(p *pending) { + b.barrier("enter", "DELIVER") + if !p.ok { + return } b.mu.Lock() - targets := make([]*subscription, 0, len(b.subs[channel])) - for _, sub := range b.subs[channel] { + targets := make([]*subscription, 0, len(b.subs[p.channel])) + for _, sub := range b.subs[p.channel] { // The shared matcher decides here and nowhere else. A second spelling is // the drift that the shared function exists to prevent. - if messaging.SubjectMatches(sub.filter, n.Subject) { + if messaging.SubjectMatches(sub.filter, p.n.Subject) { targets = append(targets, sub) } } b.mu.Unlock() for _, sub := range targets { - sub.enqueue(n.Subject, data) + sub.enqueue(p.n.Subject, p.data) } } diff --git a/core/services/pgbus/listener.go b/core/services/pgbus/listener.go index f49a58924461..62c800578d72 100644 --- a/core/services/pgbus/listener.go +++ b/core/services/pgbus/listener.go @@ -241,22 +241,78 @@ func (b *Bus) offer(channel, payload string) { } } -// resolve is where a spilled broadcast is read back and where every notification -// is dispatched. Both are off the listener on purpose, and both are on one -// goroutine, so a spilled message and an inline message of one subject keep the -// order in which they were published. +// spillFetchers is how many spilled rows one replica reads at the same time. +// +// A spilled broadcast costs one SELECT on every replica, and with one reader the +// SELECTs are in a line: a replica that reads a row in a millisecond cannot take +// more than a thousand spilled broadcasts in a second. Measured with 64 KiB +// payloads, a replica with one reader lost 10% of the broadcasts at 2,000 a +// second and 84% at 4,000. With several readers the waits overlap, and the order +// of delivery is unchanged because the dispatcher takes the notifications in the +// order they arrived. +// +// Each reader holds one connection of the pool of the replica while it reads. +const spillFetchers = 8 + +// pipelineDepth is how many notifications may be between the resolver and the +// dispatcher. It is at least spillFetchers, so that every reader has work. +const pipelineDepth = 4 * spillFetchers + +// resolve decodes the notifications that the listener copied, in the order in +// which they arrived, and starts the read of every spilled row. The dispatcher +// delivers them in the same order. // // This half may block. A notification is processed here and never on the // connection that it arrived on, so a slow SELECT here costs this replica a // bounded loss through offer and costs the server nothing. func (b *Bus) resolve() { defer close(b.resolverDone) + ordered := make(chan *pending, pipelineDepth) + dispatched := make(chan struct{}) + go b.dispatch(ordered, dispatched) + defer func() { + close(ordered) + <-dispatched + }() + + readers := make(chan struct{}, spillFetchers) for { select { case <-b.ctx.Done(): return case in := <-b.inbound: - b.deliver(in.channel, in.payload) + p := b.decode(in) + if p.ok && p.n.SpillID != "" { + select { + case readers <- struct{}{}: + case <-b.ctx.Done(): + return + } + go func() { + defer func() { <-readers }() + b.fetch(p) + }() + } + select { + case ordered <- p: + case <-b.ctx.Done(): + return + } + } + } +} + +// dispatch delivers the pending notifications in the order in which they +// arrived. A spilled one that is not ready holds the ones behind it: the order +// of the messages of one subject is kept, and the readers go on in the meantime. +func (b *Bus) dispatch(ordered <-chan *pending, done chan<- struct{}) { + defer close(done) + for p := range ordered { + select { + case <-p.ready: + b.deliver(p) + case <-b.ctx.Done(): + return } } } diff --git a/core/services/pgbus/pipeline_test.go b/core/services/pgbus/pipeline_test.go new file mode 100644 index 000000000000..996d1d38f13d --- /dev/null +++ b/core/services/pgbus/pipeline_test.go @@ -0,0 +1,149 @@ +package pgbus_test + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// A spilled broadcast costs a SELECT on every replica. The replica reads several +// rows at the same time, and delivers in the order of arrival. +var _ = Describe("reading spilled broadcasts", func() { + var ( + db *gorm.DB + pub, sub *pgbus.Bus + ) + + BeforeEach(func() { + var dsn string + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + for _, b := range []**pgbus.Bus{&pub, &sub} { + built, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, Queue: 2048, Retention: time.Hour}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(built.Close) + *b = built + } + }) + + type probe struct { + I int `json:"i"` + Fill string `json:"f"` + } + + It("delivers inline and spilled broadcasts of one subject in the order they were published", func() { + var mu sync.Mutex + var got []int + _, err := sub.Subscribe("jobs.order.stream", func(data []byte) { + var p probe + Expect(json.Unmarshal(data, &p)).To(Succeed()) + mu.Lock() + got = append(got, p.I) + mu.Unlock() + }) + Expect(err).ToNot(HaveOccurred()) + + const total = 300 + for i := 0; i < total; i++ { + fill := "x" + if i%3 == 0 { + fill = strings.Repeat("y", 16*1024) // spills + } + Expect(pub.Publish("jobs.order.stream", probe{I: i, Fill: fill})).To(Succeed()) + } + + Eventually(func() int { mu.Lock(); defer mu.Unlock(); return len(got) }, 60*time.Second).Should(Equal(total)) + for i := range got { + Expect(got[i]).To(Equal(i), "message %d arrived at position %d", got[i], i) + } + Expect(sub.Dropped()).To(BeZero()) + }) + + // The limit that was measured on a four-core database: a replica with one + // reader delivered 1,860 of 4,000 broadcasts of 64 KiB sent by 16 + // publishers, and with several readers it delivered all of them. The burst + // here is smaller, so that it holds on a busy machine too. + It("delivers a burst of 64 KiB broadcasts from many publishers without a loss", func() { + var got atomic.Int64 + _, err := sub.Subscribe("jobs.burst.big", func([]byte) { got.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + + const total = 1000 + payload := probe{Fill: strings.Repeat("a", 64*1024)} + var sent atomic.Int64 + var wg sync.WaitGroup + for w := 0; w < 16; w++ { + wg.Add(1) + go func() { + defer GinkgoRecover() + defer wg.Done() + for sent.Add(1) <= total { + Expect(pub.Publish("jobs.burst.big", payload)).To(Succeed()) + } + }() + } + wg.Wait() + + Eventually(got.Load, 60*time.Second).Should(Equal(int64(total))) + Expect(sub.Dropped()).To(BeZero()) + }) + + It("reads more than one spilled row at the same time", func() { + out := make(chan int, 64) + _, err := sub.Subscribe("jobs.overlap.probe", func(data []byte) { + var p probe + Expect(json.Unmarshal(data, &p)).To(Succeed()) + out <- p.I + }) + Expect(err).ToNot(HaveOccurred()) + + const total = 8 + for i := 0; i < total; i++ { + Expect(db.Create(&pgbus.BusMessage{ + ID: fmt.Sprintf("ov-%d", i), Subject: "jobs.overlap.probe", Payload: []byte(fmt.Sprintf(`{"i":%d}`, i)), + }).Error).To(Succeed()) + } + // A lock that stops every SELECT on the table. The statements that wait + // for it are the readers, and a replica with one reader has one. + held := db.Begin() + Expect(held.Error).ToNot(HaveOccurred()) + released := false + release := func() { + if !released { + released = true + Expect(held.Rollback().Error).ToNot(HaveOccurred()) + } + } + DeferCleanup(release) + Expect(held.Exec("LOCK TABLE bus_messages IN ACCESS EXCLUSIVE MODE").Error).To(Succeed()) + for i := 0; i < total; i++ { + Expect(db.Exec("SELECT pg_notify('localai_jobs', ?)", + fmt.Sprintf(`{"s":"jobs.overlap.probe","i":"ov-%d"}`, i)).Error).To(Succeed()) + } + + waiting := func() int64 { + var n int64 + Expect(db.Raw(`SELECT count(*) FROM pg_stat_activity + WHERE wait_event_type = 'Lock' AND query LIKE '%FROM "bus_messages"%'`).Scan(&n).Error).To(Succeed()) + return n + } + Eventually(waiting, 30*time.Second, 50*time.Millisecond).Should(BeNumerically(">", 1), + "the replica reads its spilled rows one after the other") + release() + + for want := 0; want < total; want++ { + Eventually(out, 30*time.Second).Should(Receive(Equal(want))) + } + }) +}) From ce13f6364c5305ae9d1aa2850e684b24010aed48 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 20:36:25 +0000 Subject: [PATCH 016/162] feat(carrier): build the pgbus fan-out as a member of a carrier set NewPgbusFanout opens the LISTEN connection and returns the broadcaster, the reconnect hook, a ready check and the close that a set needs. It opens the connection when it is called, so a deployment on NATS holds no LISTEN session until a change of carrier asks for one. Run the broadcaster conformance suite on holders that were swapped from one carrier to this one on a real database, and check that a message published after the swap reaches a subscriber made before it. The import-boundary spec now also keeps the PostgreSQL driver in the pgbus package, and the pgbus package behind the carrier builder. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 47 ++++- core/services/carrier/import_boundary_test.go | 76 +++++--- core/services/carrier/pgbus.go | 82 ++++++++ core/services/carrier/pgbus_test.go | 180 ++++++++++++++++++ 4 files changed, 352 insertions(+), 33 deletions(-) create mode 100644 core/services/carrier/pgbus.go create mode 100644 core/services/carrier/pgbus_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index cde386d89927..c6465927d8be 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -2,14 +2,16 @@ Distributed mode talks to other processes through a few interfaces. Code above them must not name NATS, and a new transport is a new implementation of them. -Today every seam has one carrier: NATS, or a direct dial for the gRPC and HTTP -connections to a worker. Only the carrier files import the NATS libraries -(`core/services/messaging/client.go`, `tls.go`, -`core/services/nodes/control_nats.go` and `pkg/natsauth`). A spec in -`core/services/carrier` fails when any other non-test file imports -`github.com/nats-io/*`; it reads the import lines of every Go file, so a build -constraint does not hide one. A new carrier adds its library and its files to -that spec. +Today the fan-out seam has two carriers: NATS, and `core/services/pgbus`, which +uses PostgreSQL LISTEN and NOTIFY. Every other seam has one: NATS, or a direct +dial for the gRPC and HTTP connections to a worker. Only the carrier files import +the NATS libraries (`core/services/messaging/client.go`, `tls.go`, +`core/services/nodes/control_nats.go` and `pkg/natsauth`). Only the pgbus package +imports `github.com/jackc/pgx`, and only `core/services/carrier/pgbus.go` imports +the pgbus package. A spec in `core/services/carrier` fails when any other +non-test file imports one of them; it reads the import lines of every Go file, so +a build constraint does not hide one. A new carrier adds its library and its +files to that spec. | Seam | Interface | Lives in | Today | |---|---|---|---| @@ -51,6 +53,12 @@ to `next`), then `Release(old)`. No message is published on two carriers. `NotifyReconnect` runs the reconnect hooks that `syncstate` and the failover sync register, because messages across the flip are not ordered. +`carrier.NewPgbusFanout` builds the fan-out member of a set for the pgbus +carrier. It opens the LISTEN connection when it is called and not before, so a +deployment on NATS holds no LISTEN session until a change of carrier asks for +one, and `Close` frees it. The set still needs the other members from the other +parts of that carrier before `Validate` accepts it. + The active carrier is one row of table `cluster_carrier` (`cluster.CarrierStore`). Each transition is a compare-and-set on its epoch. At startup a replica reads the row, or inserts it if it is absent, and builds the @@ -133,17 +141,38 @@ pool is keyed by host and port and two workers can report the same address. ## Rules a carrier must keep +- Payload: a broadcast above `messaging.MaxBroadcastBytes` is refused with + `messaging.ErrPayloadTooLarge` by every carrier. - Subjects: one closed set of roots, the `broadcastRoots` and `controlRoots` maps in `core/services/messaging/subject_rules.go`. Add a root there, with a row in the roots table in `subject_rules_test.go`, not in a carrier. `messaging.ValidateSubject` refuses anything else with - `messaging.ErrUnservedSubject`. + `messaging.ErrUnservedSubject`. A carrier with no request and reply, such as + pgbus, uses `messaging.ValidateBroadcastSubject`, which refuses the control + roots with the same error. `messagingtest.BroadcastRootSubjects` has one + subject for each broadcast root, and a spec keeps it equal to the rules. - Wildcards: only a whole single token `*`, never the root. `>` is refused with `messaging.ErrUnsupportedWildcard`. - Delivery is at-most-once. Anything that must survive a gap belongs in a table. - A subscriber that is reconnecting misses messages. Do not read silence as evidence about a node. +## The cluster registry + +`core/services/cluster` holds the facts that all replicas share through the +database. Table `instances` has one row for each live frontend replica +(`Registry`, kept fresh by `Membership`, which every replica runs). The +database clock stamps `last_seen`, and `instanceIsLive` is the one predicate +for "alive". The row also holds the epoch of the carrier row for which the +replica has built a carrier (`ReadyEpoch`) and the reason when it could not +(`ReadyReason`). Table `node_connections` records which live replica holds the +connection of a worker. `Registry.Claim` takes a new epoch from a sequence for +every claim, and `Release` needs both the owner and the epoch, so a replica that +lost a worker cannot clear the claim of the one that won it. Compare epochs for +equality and not for order. `Registry.Presence` tells a worker that is +reconnecting from one that is gone only when the departure is older than the +grace, and nothing but `PresenceGone` may be read as absence. + ## Four conditions that are never reported as each other 1. A routing fact: no route from here right now (`nodes.ErrNoRoute`). diff --git a/core/services/carrier/import_boundary_test.go b/core/services/carrier/import_boundary_test.go index 98f30aad6516..1b779012af01 100644 --- a/core/services/carrier/import_boundary_test.go +++ b/core/services/carrier/import_boundary_test.go @@ -27,12 +27,26 @@ import ( // Test files are exempt. They start servers and mint credentials with the // library on purpose. -const natsLibraries = "github.com/nats-io/" +// boundary is one library or package that only the files of its carrier may +// import. +type boundary struct { + // name says what the import is, for the failure message. + name string + // prefix is the import path prefix. + prefix string + // allowlist is every non-test file or package directory allowed to import + // the prefix, as a slash path from the module root. Add to it only for a + // file that is part of the carrier. + allowlist []string +} + +const ( + natsLibraries = "github.com/nats-io/" + pgxLibraries = "github.com/jackc/pgx" + pgbusPackage = "github.com/mudler/LocalAI/core/services/pgbus" +) -// natsAllowlist is every non-test file or package directory allowed to import -// natsLibraries, as a slash path from the module root. Add to it only for a -// file that is part of the NATS carrier. A new carrier library gets its own -// list in the slice that adds the carrier. +// natsAllowlist is the NATS carrier. var natsAllowlist = []string{ "core/services/messaging/client.go", "core/services/messaging/tls.go", @@ -40,6 +54,18 @@ var natsAllowlist = []string{ "pkg/natsauth/", // credential minting and decoding for the NATS carrier } +// boundaries lists what each carrier owns. A new carrier library gets its own +// row in the change that adds the carrier. +var boundaries = []boundary{ + {name: "NATS", prefix: natsLibraries, allowlist: natsAllowlist}, + // The LISTEN connection is a pgx connection and cannot come from a pool, so + // the package of the broadcast carrier is the only one that opens it. + {name: "the PostgreSQL driver", prefix: pgxLibraries, allowlist: []string{"core/services/pgbus/"}}, + // The package of the broadcast carrier is built in one place, and the code + // above the seams reaches it only through the holders. + {name: "the pgbus package", prefix: pgbusPackage, allowlist: []string{"core/services/carrier/pgbus.go"}}, +} + // importers returns the slash paths, from root, of the non-test Go files that // import a path with the given prefix. It skips hidden directories, // node_modules and nested modules. @@ -112,29 +138,31 @@ func moduleRoot() string { } } -var _ = Describe("Import boundary of the NATS libraries", func() { - It("lets only the files of the NATS carrier import them", func() { - files, err := importers(moduleRoot(), natsLibraries) - Expect(err).ToNot(HaveOccurred()) +var _ = Describe("Import boundary of the carrier libraries", func() { + for _, bd := range boundaries { + It("lets only the files of the carrier import "+bd.name, func() { + files, err := importers(moduleRoot(), bd.prefix) + Expect(err).ToNot(HaveOccurred()) - var stray []string - for _, f := range files { - if !allowed(f, natsAllowlist) { - stray = append(stray, f) + var stray []string + for _, f := range files { + if !allowed(f, bd.allowlist) { + stray = append(stray, f) + } } - } - Expect(stray).To(BeEmpty(), "these files import %s but are not in natsAllowlist; code above the seams must use messaging and nodes interfaces", natsLibraries) - }) + Expect(stray).To(BeEmpty(), "these files import %s but are not in the allowlist; code above the seams must use messaging and nodes interfaces", bd.prefix) + }) - It("keeps the allowlist honest: every entry still imports the libraries", func() { - files, err := importers(moduleRoot(), natsLibraries) - Expect(err).ToNot(HaveOccurred()) + It("keeps the allowlist of "+bd.name+" honest: every entry still imports it", func() { + files, err := importers(moduleRoot(), bd.prefix) + Expect(err).ToNot(HaveOccurred()) - for _, a := range natsAllowlist { - used := slices.ContainsFunc(files, func(f string) bool { return allowed(f, []string{a}) }) - Expect(used).To(BeTrue(), "%s no longer imports %s; remove it from natsAllowlist", a, natsLibraries) - } - }) + for _, a := range bd.allowlist { + used := slices.ContainsFunc(files, func(f string) bool { return allowed(f, []string{a}) }) + Expect(used).To(BeTrue(), "%s no longer imports %s; remove it from the allowlist", a, bd.prefix) + } + }) + } Describe("the scan", func() { var root string diff --git a/core/services/carrier/pgbus.go b/core/services/carrier/pgbus.go new file mode 100644 index 000000000000..5da79e0327be --- /dev/null +++ b/core/services/carrier/pgbus.go @@ -0,0 +1,82 @@ +package carrier + +import ( + "context" + "errors" + "fmt" + + "go.opentelemetry.io/otel/metric" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/advisorylock" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/pgbus" +) + +// PgbusOptions is what NewPgbusFanout needs. +type PgbusOptions struct { + // DB is the pooled handle of the cluster database. + DB *gorm.DB + // DSN is the connection string of the LISTEN connection. It is the string of + // the database that DB points to, because a LISTEN connection cannot come + // from a pool. + DSN string + // Meter receives the counters of the carrier. Nil means the global meter. + Meter metric.Meter +} + +// Fanout is the fan-out member of a Set for a carrier that has no NATS server: +// the broadcaster, the reconnect hook the holder reads, a ready check and the +// close. The other members of the set come from the other parts of the carrier. +type Fanout struct { + // Broadcaster delivers through LISTEN and NOTIFY. + Broadcaster messaging.Broadcaster + // OnReconnect registers a callback that runs after the LISTEN connection was + // found again. It is the OnReconnect of the Set. + OnReconnect func(func()) + // Ready reports an error when the LISTEN connection is down. The report that + // a replica sends for a change of carrier reads it. A caller must not decide + // on it anything about a worker. + Ready func() error + // Dropped reports how many broadcasts this replica received and lost. + Dropped func() uint64 + // Close ends the LISTEN connection. It is the Close of the Set. + Close func() +} + +// NewPgbusFanout opens the LISTEN connection and returns the fan-out member. +// +// It is the one place that opens the connection. A deployment whose active +// carrier is NATS does not call it until a change of carrier asks for it, so it +// holds no LISTEN session, and it closes the member after the drain, so it holds +// none again. A LISTEN session needs a direct connection to the database, or a +// pooler in session mode. +func NewPgbusFanout(ctx context.Context, o PgbusOptions) (*Fanout, error) { + if o.DB == nil { + return nil, errors.New("pgbus fan-out needs a database handle") + } + if o.DSN == "" { + return nil, errors.New("pgbus fan-out needs the connection string of the database") + } + if err := advisorylock.WithLockCtx(ctx, o.DB, advisorylock.KeySchemaMigrate, func() error { + return pgbus.Migrate(ctx, o.DB) + }); err != nil { + return nil, fmt.Errorf("migrating the spill table: %w", err) + } + bus, err := pgbus.New(ctx, pgbus.Config{DSN: o.DSN, DB: o.DB, Meter: o.Meter}) + if err != nil { + return nil, err + } + return &Fanout{ + Broadcaster: bus, + OnReconnect: bus.OnReconnect, + Ready: func() error { + if !bus.IsConnected() { + return errors.New("the LISTEN connection to the database is down") + } + return nil + }, + Dropped: bus.Dropped, + Close: bus.Close, + }, nil +} diff --git a/core/services/carrier/pgbus_test.go b/core/services/carrier/pgbus_test.go new file mode 100644 index 000000000000..186778bb734b --- /dev/null +++ b/core/services/carrier/pgbus_test.go @@ -0,0 +1,180 @@ +package carrier_test + +import ( + "context" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// listeners counts the LISTEN sessions that the pgbus carrier holds on the +// database. +func listeners(db *gorm.DB) int64 { + GinkgoHelper() + var n int64 + Expect(db.Raw("SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'").Scan(&n).Error).To(Succeed()) + return n +} + +// pgbusSet is a set of the tunnel carrier with the pgbus fan-out. The other +// members are fakes: later changes provide them. +func pgbusSet(ctx context.Context, db *gorm.DB, dsn string, epoch int64) (*carrier.Set, *carrier.Fanout) { + GinkgoHelper() + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(fan.Close) + set := newFakeCarrier(cluster.CarrierTunnel, epoch).set + set.Broadcaster = fan.Broadcaster + set.OnReconnect = fan.OnReconnect + set.Close = fan.Close + Expect(set.Validate()).To(Succeed()) + return set, fan +} + +var _ = Describe("The pgbus fan-out of the tunnel carrier", func() { + var ( + db *gorm.DB + dsn string + ctx context.Context + ) + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + ctx = context.Background() + }) + + It("opens no connection until it is built, and none after it is closed", func() { + // A deployment on NATS builds nothing, so it holds no LISTEN session. + Expect(listeners(db)).To(BeZero()) + + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) + Expect(err).ToNot(HaveOccurred()) + Expect(listeners(db)).To(Equal(int64(1))) + + fan.Close() + Eventually(func() int64 { return listeners(db) }, 10*time.Second).Should(BeZero()) + }) + + It("reports whether the LISTEN connection is up, for the ready report", func() { + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) + Expect(err).ToNot(HaveOccurred()) + Expect(fan.Ready()).To(Succeed()) + + fan.Close() + Expect(fan.Ready()).To(HaveOccurred()) + }) + + It("creates the spill table when it is built", func() { + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(fan.Close) + + Expect(db.Migrator().HasTable("bus_messages")).To(BeTrue()) + }) + + It("refuses to be built without a database string, and holds no session then", func() { + _, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db}) + Expect(err).To(HaveOccurred()) + Expect(listeners(db)).To(BeZero()) + }) + + It("builds a valid set of the tunnel carrier", func() { + set, _ := pgbusSet(ctx, db, dsn, 4) + Expect(set.Name).To(Equal(cluster.CarrierTunnel)) + Expect(set.Epoch).To(Equal(int64(4))) + }) + + Describe("behind the holder", func() { + It("receives, after a swap, what a peer publishes, and keeps what was subscribed before", func() { + // Two replicas. Each starts on its own fake NATS carrier, with a + // subscription made before the swap. Each swaps to a pgbus set on + // the same database, in the order the switch uses. + type replica struct { + holder *carrier.Broadcaster + cur *atomic.Pointer[carrier.Set] + first *carrier.Set + next *carrier.Set + got chan []byte + } + build := func() *replica { + r := &replica{cur: &atomic.Pointer[carrier.Set]{}, got: make(chan []byte, 8)} + r.first = newFakeCarrier(cluster.CarrierNATS, 1).set + r.cur.Store(r.first) + r.holder = carrier.NewBroadcaster(r.cur) + _, err := r.holder.Subscribe("jobs.j1.result", func(b []byte) { r.got <- b }) + Expect(err).ToNot(HaveOccurred()) + r.next, _ = pgbusSet(ctx, db, dsn, 2) + return r + } + a, b := build(), build() + + // prepare: both listen on both carriers. + Expect(a.holder.Listen(a.next)).To(Succeed()) + Expect(b.holder.Listen(b.next)).To(Succeed()) + // commit: a publishes on the new carrier, b has not seen it yet. + a.cur.Store(a.next) + Expect(a.holder.Publish("jobs.j1.result", "from-a")).To(Succeed()) + Eventually(b.got, 10*time.Second).Should(Receive(Equal([]byte(`"from-a"`)))) + Eventually(a.got, 10*time.Second).Should(Receive(Equal([]byte(`"from-a"`)))) + + // b flips and publishes on the new carrier too. a still hears it. + b.cur.Store(b.next) + Expect(b.holder.Publish("jobs.j1.result", "from-b")).To(Succeed()) + Eventually(a.got, 10*time.Second).Should(Receive(Equal([]byte(`"from-b"`)))) + + // stable: the old carriers are released and the subscription lives on. + Expect(a.holder.Release(a.first)).To(Succeed()) + Expect(b.holder.Release(b.first)).To(Succeed()) + Expect(a.holder.Publish("jobs.j1.result", "after")).To(Succeed()) + Eventually(b.got, 10*time.Second).Should(Receive(Equal([]byte(`"after"`)))) + }) + + It("runs the reconnect hooks of the holder when the LISTEN connection is lost and found again", func() { + cur := &atomic.Pointer[carrier.Set]{} + set, _ := pgbusSet(ctx, db, dsn, 1) + cur.Store(set) + holder := carrier.NewBroadcaster(cur) + called := make(chan struct{}, 4) + holder.OnReconnect(func() { called <- struct{}{} }) + _, err := holder.Subscribe("jobs.j2.result", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + + Expect(db.Exec("SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'").Error).To(Succeed()) + + Eventually(called, 30*time.Second).Should(Receive()) + }) + }) + + // The suite that every carrier passes, run on holders that were swapped from + // a NATS set to a pgbus set, over two replicas. + Describe("Broadcaster conformance on holders swapped to pgbus", func() { + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { + pgDB, pgDSN := testutil.SetupTestDBWithDSN() + swapped := func() (*carrier.Broadcaster, *carrier.Fanout) { + var cur atomic.Pointer[carrier.Set] + first := newFakeCarrier(cluster.CarrierNATS, 1).set + cur.Store(first) + h := carrier.NewBroadcaster(&cur) + next, fan := pgbusSet(context.Background(), pgDB, pgDSN, 2) + Expect(h.Listen(next)).To(Succeed()) + cur.Store(next) + Expect(h.Release(first)).To(Succeed()) + return h, fan + } + h1, f1 := swapped() + h2, f2 := swapped() + return messagingtest.Carrier{ + Bus: h1, Peer: h2, + Dropped: func() uint64 { return f1.Dropped() + f2.Dropped() }, + } + }) + }) +}) From ae780e75d8dd7e4d9d6b925aeb1774767dc29e6f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:25:14 +0000 Subject: [PATCH 017/162] feat(tunnel): add the wire format and the session of the worker tunnel A worker holds one outbound websocket to a frontend and a yamux session runs on it. The frontend opens a stream for every request. The first frame of a stream names the local service of the worker, and the worker always answers with a reply frame before the carried protocol starts. The four refusals stay apart, so that a caller acts on the three that are evidence about a backend and does nothing about the one that says the worker learned nothing. The upgrader and the dialer come from one constructor each, with 64 KiB buffers. With the 4 KiB default of gorilla, a stream that sends from the worker to the frontend ran at 560 to 600 MB/s on a loopback socket and the other direction at 3,200 to 3,800 MB/s. With 64 KiB the first direction runs at 2,900 to 3,000 MB/s. BenchmarkTransfer repeats the measurement. A session has a lane. The inference lane keeps the yamux windows. The bulk lane uses 4 MiB and 32 MiB, which a transfer needs on a link with a long round trip, and which would let more data queue ahead of a small call on a session that is shared. Only the files of the tunnel package import yamux. The import boundary spec lists it. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/import_boundary_test.go | 5 + core/services/tunnel/bench_test.go | 105 +++++++ core/services/tunnel/export_test.go | 6 + core/services/tunnel/helpers_test.go | 67 ++++ core/services/tunnel/proto.go | 292 ++++++++++++++++++ core/services/tunnel/proto_test.go | 134 ++++++++ core/services/tunnel/session.go | 128 ++++++++ core/services/tunnel/session_test.go | 154 +++++++++ core/services/tunnel/tunnel_suite_test.go | 13 + core/services/tunnel/ws.go | 35 +++ core/services/tunnel/wsconn.go | 134 ++++++++ go.mod | 2 +- 12 files changed, 1074 insertions(+), 1 deletion(-) create mode 100644 core/services/tunnel/bench_test.go create mode 100644 core/services/tunnel/export_test.go create mode 100644 core/services/tunnel/helpers_test.go create mode 100644 core/services/tunnel/proto.go create mode 100644 core/services/tunnel/proto_test.go create mode 100644 core/services/tunnel/session.go create mode 100644 core/services/tunnel/session_test.go create mode 100644 core/services/tunnel/tunnel_suite_test.go create mode 100644 core/services/tunnel/ws.go create mode 100644 core/services/tunnel/wsconn.go diff --git a/core/services/carrier/import_boundary_test.go b/core/services/carrier/import_boundary_test.go index 1b779012af01..56d6396277a9 100644 --- a/core/services/carrier/import_boundary_test.go +++ b/core/services/carrier/import_boundary_test.go @@ -44,6 +44,7 @@ const ( natsLibraries = "github.com/nats-io/" pgxLibraries = "github.com/jackc/pgx" pgbusPackage = "github.com/mudler/LocalAI/core/services/pgbus" + yamuxLibrary = "github.com/libp2p/go-yamux/" ) // natsAllowlist is the NATS carrier. @@ -64,6 +65,10 @@ var boundaries = []boundary{ // The package of the broadcast carrier is built in one place, and the code // above the seams reaches it only through the holders. {name: "the pgbus package", prefix: pgbusPackage, allowlist: []string{"core/services/carrier/pgbus.go"}}, + // The multiplexer of the tunnel carrier. The connect endpoint and the + // worker hold a tunnel session, but they reach it through the types of the + // tunnel package. + {name: "yamux", prefix: yamuxLibrary, allowlist: []string{"core/services/tunnel/"}}, } // importers returns the slash paths, from root, of the non-test Go files that diff --git a/core/services/tunnel/bench_test.go b/core/services/tunnel/bench_test.go new file mode 100644 index 000000000000..ec3e40d697c2 --- /dev/null +++ b/core/services/tunnel/bench_test.go @@ -0,0 +1,105 @@ +package tunnel_test + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/mudler/LocalAI/core/services/tunnel" +) + +// benchPair is sessionPair for benchmarks, with the upgrader and the dialer +// given by the caller so that a run can use other buffer sizes. +func benchPair(b *testing.B, up *websocket.Upgrader, dialer *websocket.Dialer, lane tunnel.Lane) (frontend, worker *tunnel.Session) { + b.Helper() + ch := make(chan *tunnel.Session, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + s, err := tunnel.ServerSession(ws, lane) + if err != nil { + _ = ws.Close() + return + } + ch <- s + })) + b.Cleanup(srv.Close) + ws, _, err := dialer.Dial("ws"+srv.URL[len("http"):], nil) + if err != nil { + b.Fatal(err) + } + worker, err = tunnel.ClientSession(ws, lane) + if err != nil { + b.Fatal(err) + } + b.Cleanup(func() { _ = worker.Close() }) + frontend = <-ch + b.Cleanup(func() { _ = frontend.Close() }) + return frontend, worker +} + +// BenchmarkTransfer measures one stream that moves 64 MiB between a frontend +// and a worker over a loopback socket, with the websocket buffers of gorilla +// (4 KiB) and the buffers of the tunnel (64 KiB). Push is the worker sending to +// the frontend, which is the direction where the sender masks the payload. +// +// go test ./core/services/tunnel -run '^$' -bench Transfer -benchtime 5x +func BenchmarkTransfer(b *testing.B) { + const size = 64 << 20 + buffers := []struct { + name string + up *websocket.Upgrader + dial *websocket.Dialer + }{ + {"buffer4KiB", &websocket.Upgrader{}, &websocket.Dialer{HandshakeTimeout: 5 * time.Second}}, + {"buffer64KiB", tunnel.NewUpgrader(), tunnel.NewDialer(5 * time.Second)}, + } + for _, buf := range buffers { + for _, dir := range []string{"push", "pull"} { + b.Run(dir+"/"+buf.name, func(b *testing.B) { + frontend, worker := benchPair(b, buf.up, buf.dial, tunnel.LaneInference) + sender, receiver := worker, frontend + if dir == "pull" { + sender, receiver = frontend, worker + } + chunk := make([]byte, 256<<10) + b.SetBytes(size) + b.ResetTimer() + for range b.N { + done := make(chan error, 1) + go func() { + st, err := receiver.AcceptStream() + if err != nil { + done <- err + return + } + _, err = io.Copy(io.Discard, st) + _ = st.Close() + done <- err + }() + st, err := sender.OpenStream(context.Background()) + if err != nil { + b.Fatal(err) + } + for sent := 0; sent < size; sent += len(chunk) { + if _, err := st.Write(chunk); err != nil { + b.Fatal(err) + } + } + if err := st.Close(); err != nil { + b.Fatal(err) + } + if err := <-done; err != nil { + b.Fatal(err) + } + } + }) + } + } +} diff --git a/core/services/tunnel/export_test.go b/core/services/tunnel/export_test.go new file mode 100644 index 000000000000..611cc6b36776 --- /dev/null +++ b/core/services/tunnel/export_test.go @@ -0,0 +1,6 @@ +package tunnel + +import "github.com/libp2p/go-yamux/v5" + +// SessionConfigFor exposes the yamux configuration of a lane to the specs. +func SessionConfigFor(lane Lane) *yamux.Config { return sessionConfig(lane) } diff --git a/core/services/tunnel/helpers_test.go b/core/services/tunnel/helpers_test.go new file mode 100644 index 000000000000..8a41c490f8f1 --- /dev/null +++ b/core/services/tunnel/helpers_test.go @@ -0,0 +1,67 @@ +package tunnel_test + +import ( + "net" + "net/http" + "net/http/httptest" + "time" + + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// sessionPair opens a real websocket on a loopback socket and starts the two +// sides of a session on it. The server side is the frontend and the client side +// is the worker. +func sessionPair(lane tunnel.Lane) (server, client *tunnel.Session) { + GinkgoHelper() + upgrader := tunnel.NewUpgrader() + serverCh := make(chan *tunnel.Session, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + s, err := tunnel.ServerSession(ws, lane) + if err != nil { + _ = ws.Close() + return + } + serverCh <- s + })) + DeferCleanup(srv.Close) + + url := "ws" + srv.URL[len("http"):] + ws, _, err := tunnel.NewDialer(5*time.Second).Dial(url, nil) + Expect(err).ToNot(HaveOccurred()) + client, err = tunnel.ClientSession(ws, lane) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = client.Close() }) + + Eventually(serverCh).Should(Receive(&server)) + DeferCleanup(func() { _ = server.Close() }) + return server, client +} + +// tcpPair returns the two ends of a loopback TCP connection. +func tcpPair() (a, b *net.TCPConn) { + GinkgoHelper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = lis.Close() }() + accepted := make(chan *net.TCPConn, 1) + go func() { + c, err := lis.Accept() + if err == nil { + accepted <- c.(*net.TCPConn) + } + }() + c, err := net.Dial("tcp", lis.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + a = c.(*net.TCPConn) + Eventually(accepted).Should(Receive(&b)) + DeferCleanup(func() { _ = a.Close(); _ = b.Close() }) + return a, b +} diff --git a/core/services/tunnel/proto.go b/core/services/tunnel/proto.go new file mode 100644 index 000000000000..d53d543ff3e9 --- /dev/null +++ b/core/services/tunnel/proto.go @@ -0,0 +1,292 @@ +// Package tunnel carries the traffic between frontends and a worker over one +// outbound connection of the worker. +// +// The worker dials a frontend and opens a websocket. A yamux session runs on +// that websocket. The frontend opens a stream for every request and the worker +// answers it, so the worker needs no inbound port. This package holds the wire +// format, the session, the registry of the sessions a frontend holds, and the +// helpers that join two streams. It is the only package, with its callers in the +// connect endpoint and the worker, that imports the multiplexer. +package tunnel + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + "strings" + "unicode/utf8" +) + +// The first thing on every stream is a request frame. A yamux stream carries no +// destination, so the frame names a tag (which local service of the worker) and +// a target (which instance of it). The worker answers with a reply frame before +// either side speaks the carried protocol. +// +// The reply is always sent, not only on failure. The carried protocols speak +// from the client side first (gRPC sends a preface, HTTP sends a request line), +// so a reply sent only on failure would arrive between the bytes of a response +// on the streams that worked, and the frontend would not know when to read it. +// +// A frame has a 2-byte length. The reader takes exactly the frame and nothing +// after it, because the stream goes to gRPC or net/http next and a reader that +// buffered too much would take the first bytes of their conversation. + +const ( + // StreamTagGRPC sends a stream to a backend process of the worker. The + // target is the address of the backend. The worker keeps only the port. + StreamTagGRPC = "grpc" + + // StreamTagHTTP sends a stream to the HTTP server of the worker, which + // serves file staging and backend logs. The target is ignored: a worker has + // one such server and only the worker knows where it listens. + StreamTagHTTP = "http" +) + +// maxFrame bounds a frame. The longest real frame is a tag and an address, under +// a hundred bytes. The bound stops a peer that declares a large length and then +// sends nothing from making the reader allocate. +const maxFrame = 1024 + +// The reply codes go on the wire as strings, so a frontend that reads a code it +// does not know can still log something an operator can search for. +const ( + replyAccepted = "ok" + replyCodeUnknownTag = "unknown-tag" + replyCodeUnavailable = "unavailable" + replyCodeBadRequest = "bad-request" + replyCodeNotServed = "not-served" + replyPrefixRefused = "err " + streamRequestSeparator = " " +) + +// The four refusals a worker can send. They stay apart because a caller acts on +// each in a different way. A caller gives up on an unknown tag, because it does +// not change until the worker is upgraded. It acts on an unavailable target, +// which is what a backend that died looks like from inside the worker. It +// reports a bad request, which is a bug of the frontend. It does nothing about +// "not served", which says that the worker learned nothing. +// +// None of them wraps an error about absence of the node. A refusal proves that +// the worker is connected and that it answered. +var ( + ErrStreamTagUnknown = errors.New("tunnel: the worker does not serve that stream tag") + ErrStreamTargetUnavailable = errors.New("tunnel: the worker could not reach the local service for that stream") + ErrStreamRequestInvalid = errors.New("tunnel: the worker rejected the stream request as malformed") + + // ErrStreamNotServed means that the worker could not serve the stream for a + // reason of its own, and that the reason says nothing about the backend that + // the stream named. IsWorkerAnswer does not count it as evidence, so it + // reaches a consumer as a routing failure and nothing is reaped. + // + // A frontend that does not know this code reads it as an unrecognised reply, + // which ReadStreamReply returns as a plain error. A mixed deployment is + // therefore safe without extra work. + ErrStreamNotServed = errors.New("tunnel: the worker could not serve that stream, for a reason that is not about the backend") +) + +// streamRefusals is the whole refusal vocabulary in one table. The writer, the +// reader, IsWorkerAnswer and IsStreamRefusal all read it, so a new refusal +// cannot be added to some of them and forgotten in the others. +// +// The writer takes the first entry that matches with errors.Is, so a reason that +// wraps two sentinels gives the same code every time. +// +// evidence says if a consumer may act on the code. Three codes are statements +// about a backend. ErrStreamNotServed is the worker saying that it learned +// nothing, and counting it as evidence would turn every worker-side failure that +// is temporary into an eviction. +var streamRefusals = []struct { + sentinel error + code string + evidence bool +}{ + {ErrStreamTagUnknown, replyCodeUnknownTag, true}, + {ErrStreamTargetUnavailable, replyCodeUnavailable, true}, + {ErrStreamRequestInvalid, replyCodeBadRequest, true}, + {ErrStreamNotServed, replyCodeNotServed, false}, +} + +// IsWorkerAnswer reports whether err is a refusal that a consumer may act on. +// It is true for the three codes that are statements about a backend and false +// for everything else, including ErrStreamNotServed, a failure to read the +// reply and an error of the transport. +func IsWorkerAnswer(err error) bool { + for _, r := range streamRefusals { + if r.evidence && errors.Is(err, r.sentinel) { + return true + } + } + return false +} + +// IsStreamRefusal reports whether err already carries a classification of this +// vocabulary. It is a different question from IsWorkerAnswer. A worker that +// classifies an error again would overwrite a decision that was made closer to +// the failure. +func IsStreamRefusal(err error) bool { + for _, r := range streamRefusals { + if errors.Is(err, r.sentinel) { + return true + } + } + return false +} + +// WriteStreamRequest sends the opening frame that names what the stream is for. +// An empty tag, or a tag with a space, is refused here. The worker would refuse +// it too, but the caller learns a round trip later. +func WriteStreamRequest(w io.Writer, tag, target string) error { + if tag == "" { + return errors.New("writing a tunnel stream request: empty tag") + } + if strings.Contains(tag, streamRequestSeparator) { + // The split is on the first space, so a space in the tag would move a + // part of the tag into the target. + return fmt.Errorf("writing a tunnel stream request: tag %q contains a space", tag) + } + return writeFrame(w, tag+streamRequestSeparator+target) +} + +// ReadStreamRequest reads the opening frame. The target is empty when the tag +// has no argument. +// +// A malformed frame is a plain error and not ErrStreamRequestInvalid. That +// sentinel is what a worker sends to describe a refusal, and a caller must be +// able to tell "the peer refused my request" from "I could not read the request +// of the peer". +func ReadStreamRequest(r io.Reader) (tag, target string, err error) { + payload, err := readFrame(r) + if err != nil { + return "", "", fmt.Errorf("reading a tunnel stream request: %w", err) + } + tag, target, _ = strings.Cut(payload, streamRequestSeparator) + if tag == "" { + return "", "", errors.New("reading a tunnel stream request: empty tag") + } + return tag, target, nil +} + +// WriteStreamAccepted tells the frontend that the stream now carries the +// protocol. Everything after this frame belongs to that protocol. +func WriteStreamAccepted(w io.Writer) error { + return writeFrame(w, replyAccepted) +} + +// WriteStreamRefusal says why a stream is not served. The caller closes the +// stream after it. +// +// A reason with no known classification is sent as "not served" with its text, +// and not as a bad request. The other three codes are evidence that a frontend +// acts on, up to the removal of a model row. An error that arrives here with no +// classification was not classified by anyone, and it must not become a verdict +// by default. +func WriteStreamRefusal(w io.Writer, reason error) error { + code := replyCodeNotServed + for _, r := range streamRefusals { + if errors.Is(reason, r.sentinel) { + code = r.code + break + } + } + + text := "" + if reason != nil { + text = strings.Map(func(r rune) rune { + // A new line would not break the frame, which has a length. The + // text goes to a log line on the other side, and a cause on several + // lines is hard to search. + if r == '\n' || r == '\r' { + return ' ' + } + return r + }, reason.Error()) + } + frame := replyPrefixRefused + code + streamRequestSeparator + text + return writeFrame(w, truncateRunes(frame, maxFrame)) +} + +// ReadStreamReply reads the answer of the worker. A nil result means that the +// stream now carries the protocol. +// +// A failure to read the reply is returned as itself and never as a refusal. A +// refusal means that the worker is connected and said no. A failed read means +// that the tunnel broke, and a caller that took one for the other would report +// a dead link as a decision of the worker. +func ReadStreamReply(r io.Reader) error { + payload, err := readFrame(r) + if err != nil { + return fmt.Errorf("reading a tunnel stream reply: %w", err) + } + if payload == replyAccepted { + return nil + } + rest, ok := strings.CutPrefix(payload, replyPrefixRefused) + if !ok { + return fmt.Errorf("reading a tunnel stream reply: unrecognised reply %q", payload) + } + code, text, _ := strings.Cut(rest, streamRequestSeparator) + for _, r := range streamRefusals { + if code == r.code { + return fmt.Errorf("%w: %s", r.sentinel, text) + } + } + // A code from a newer worker. It is not mapped onto the nearest known code, + // so that a frontend does not retry for ever against a refusal with another + // meaning, and so that IsWorkerAnswer is false and nothing is reaped. + return fmt.Errorf("tunnel stream refused with unrecognised code %q: %s", code, text) +} + +// truncateRunes cuts s to at most limit bytes on a rune boundary. A plain slice +// could cut a rune in two and put a lone continuation byte on the wire. +func truncateRunes(s string, limit int) string { + if len(s) <= limit { + return s + } + cut := limit + // A rune has at most 4 bytes, so this walks back at most 3 steps. + for cut > 0 && !utf8.RuneStart(s[cut]) { + cut-- + } + return s[:cut] +} + +// writeFrame writes a frame with one Write. The stream is a yamux stream, where +// a write becomes one data frame, and two writes would put the length and the +// payload in two frames. One Write also keeps the websocket adapter to one +// message for each frame. +func writeFrame(w io.Writer, payload string) error { + if len(payload) > maxFrame { + return fmt.Errorf("tunnel frame is %d bytes, over the %d-byte limit", len(payload), maxFrame) + } + buf := make([]byte, 2+len(payload)) + // #nosec G115 -- the check above limits the payload to maxFrame (1024), + // which is inside uint16. + binary.BigEndian.PutUint16(buf[:2], uint16(len(payload))) + copy(buf[2:], payload) + _, err := w.Write(buf) + return err +} + +// readFrame reads a frame. It uses io.ReadFull because a yamux stream returns +// what has arrived, and a header in two data frames is normal. A frame that is +// cut short gives io.ErrUnexpectedEOF, which is what a peer that hung up in the +// middle of a header looks like. +func readFrame(r io.Reader) (string, error) { + var size [2]byte + if _, err := io.ReadFull(r, size[:]); err != nil { + return "", err + } + n := binary.BigEndian.Uint16(size[:]) + if int(n) > maxFrame { + return "", fmt.Errorf("tunnel frame declares %d bytes, over the %d-byte limit", n, maxFrame) + } + if n == 0 { + return "", nil + } + payload := make([]byte, n) + if _, err := io.ReadFull(r, payload); err != nil { + return "", err + } + return string(payload), nil +} diff --git a/core/services/tunnel/proto_test.go b/core/services/tunnel/proto_test.go new file mode 100644 index 000000000000..d8af3de2d6c2 --- /dev/null +++ b/core/services/tunnel/proto_test.go @@ -0,0 +1,134 @@ +package tunnel_test + +import ( + "bytes" + "errors" + "fmt" + "io" + "strings" + + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Stream protocol", func() { + It("carries a request frame without reading a byte of what follows", func() { + var buf bytes.Buffer + Expect(tunnel.WriteStreamRequest(&buf, tunnel.StreamTagGRPC, "127.0.0.1:50052")).To(Succeed()) + buf.WriteString("PRI * HTTP/2.0") + + tag, target, err := tunnel.ReadStreamRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(tag).To(Equal("grpc")) + Expect(target).To(Equal("127.0.0.1:50052")) + Expect(buf.String()).To(Equal("PRI * HTTP/2.0")) + }) + + It("accepts a tag with no target", func() { + var buf bytes.Buffer + Expect(tunnel.WriteStreamRequest(&buf, tunnel.StreamTagHTTP, "")).To(Succeed()) + tag, target, err := tunnel.ReadStreamRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(tag).To(Equal("http")) + Expect(target).To(BeEmpty()) + }) + + It("refuses to write an empty tag or a tag with a space", func() { + Expect(tunnel.WriteStreamRequest(io.Discard, "", "x")).ToNot(Succeed()) + Expect(tunnel.WriteStreamRequest(io.Discard, "gr pc", "x")).ToNot(Succeed()) + }) + + It("does not return a malformed request as a refusal sentinel", func() { + _, _, err := tunnel.ReadStreamRequest(bytes.NewReader([]byte{0, 0})) + Expect(err).To(HaveOccurred()) + Expect(tunnel.IsStreamRefusal(err)).To(BeFalse()) + }) + + It("refuses a frame that declares more than the limit before it reads the payload", func() { + r := bytes.NewReader([]byte{0xff, 0xff, 'x'}) + _, _, err := tunnel.ReadStreamRequest(r) + Expect(err).To(MatchError(ContainSubstring("over the"))) + }) + + It("reads a frame that is cut short as an unexpected end", func() { + _, _, err := tunnel.ReadStreamRequest(bytes.NewReader([]byte{0, 9, 'a'})) + Expect(err).To(MatchError(io.ErrUnexpectedEOF)) + }) + + It("reads an accepted reply as success", func() { + var buf bytes.Buffer + Expect(tunnel.WriteStreamAccepted(&buf)).To(Succeed()) + Expect(tunnel.ReadStreamReply(&buf)).To(Succeed()) + }) + + DescribeTable("round-trips every refusal to its sentinel", + func(sentinel error, evidence bool) { + var buf bytes.Buffer + Expect(tunnel.WriteStreamRefusal(&buf, fmt.Errorf("%w: because", sentinel))).To(Succeed()) + err := tunnel.ReadStreamReply(&buf) + Expect(err).To(MatchError(sentinel)) + Expect(tunnel.IsStreamRefusal(err)).To(BeTrue()) + Expect(tunnel.IsWorkerAnswer(err)).To(Equal(evidence)) + }, + Entry("unknown tag", tunnel.ErrStreamTagUnknown, true), + Entry("target unavailable", tunnel.ErrStreamTargetUnavailable, true), + Entry("request invalid", tunnel.ErrStreamRequestInvalid, true), + Entry("not served", tunnel.ErrStreamNotServed, false), + ) + + It("sends a reason with no classification as not served, never as evidence", func() { + var buf bytes.Buffer + Expect(tunnel.WriteStreamRefusal(&buf, errors.New("something unclassified"))).To(Succeed()) + err := tunnel.ReadStreamReply(&buf) + Expect(err).To(MatchError(tunnel.ErrStreamNotServed)) + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("keeps the first sentinel of the table when a reason wraps two", func() { + both := errors.Join(tunnel.ErrStreamNotServed, tunnel.ErrStreamTagUnknown) + var buf bytes.Buffer + Expect(tunnel.WriteStreamRefusal(&buf, both)).To(Succeed()) + Expect(tunnel.ReadStreamReply(&buf)).To(MatchError(tunnel.ErrStreamTagUnknown)) + }) + + It("reads a code it does not know as an error that is no evidence", func() { + var buf bytes.Buffer + Expect(tunnel.WriteStreamRefusal(&buf, errors.New("x"))).To(Succeed()) + raw := strings.Replace(buf.String(), "not-served", "from-the-future", 1) + // The length prefix covers the old text, so build the frame again. + var framed bytes.Buffer + framed.Write([]byte{0, byte(len(raw) - 2)}) + framed.WriteString(raw[2:]) + + err := tunnel.ReadStreamReply(&framed) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("from-the-future")) + Expect(tunnel.IsStreamRefusal(err)).To(BeFalse()) + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("does not report a failed read of the reply as a refusal", func() { + err := tunnel.ReadStreamReply(bytes.NewReader(nil)) + Expect(err).To(HaveOccurred()) + Expect(tunnel.IsStreamRefusal(err)).To(BeFalse()) + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("cuts a long reason on a rune boundary and puts it on one line", func() { + reason := errors.New(strings.Repeat("é\n", 2000)) + var buf bytes.Buffer + Expect(tunnel.WriteStreamRefusal(&buf, reason)).To(Succeed()) + Expect(buf.Len()).To(BeNumerically("<=", 2+1024)) + err := tunnel.ReadStreamReply(&buf) + Expect(err).To(MatchError(tunnel.ErrStreamNotServed)) + Expect(err.Error()).ToNot(ContainSubstring("\n")) + Expect(err.Error()).To(MatchRegexp(`^[^\x{FFFD}]*$`)) + }) + + It("does not treat a plain error or nil as an answer of a worker", func() { + Expect(tunnel.IsWorkerAnswer(nil)).To(BeFalse()) + Expect(tunnel.IsWorkerAnswer(errors.New("dial"))).To(BeFalse()) + }) +}) diff --git a/core/services/tunnel/session.go b/core/services/tunnel/session.go new file mode 100644 index 000000000000..91772ae87967 --- /dev/null +++ b/core/services/tunnel/session.go @@ -0,0 +1,128 @@ +package tunnel + +import ( + "context" + "fmt" + "io" + "net" + + "github.com/gorilla/websocket" + "github.com/libp2p/go-yamux/v5" +) + +// Lane names one of the two sessions that a worker holds to a frontend. +// +// The inference lane carries the requests that need a short delay: model +// calls, control and health. The bulk lane carries file transfers. A large +// transfer on the inference lane delays every other stream of the session by the +// amount of data that is in flight ahead of them, because yamux queues up to one +// window of data for each stream before a new frame gets its turn on the +// websocket. A separate session has its own websocket and its own window, so the +// two cannot delay each other. +type Lane string + +const ( + LaneInference Lane = "inference" + LaneBulk Lane = "bulk" +) + +// Valid reports whether l is a lane that this package knows. +func (l Lane) Valid() bool { return l == LaneInference || l == LaneBulk } + +// ParseLane maps the value of the lane parameter of the connect request to a +// lane. The empty value is the inference lane, which is what a worker that +// predates the bulk lane connects to. +func ParseLane(s string) (Lane, error) { + if s == "" { + return LaneInference, nil + } + if l := Lane(s); l.Valid() { + return l, nil + } + return "", fmt.Errorf("unknown tunnel lane %q", s) +} + +// The bulk windows are larger than the defaults of yamux (256 KiB at the start, +// 16 MiB at most). With the defaults a transfer on a link with 20 ms of latency +// runs about 20% slower. The larger windows also let more data queue ahead of a +// new frame, so they would raise the delay of small calls on a shared session. +// That is why they apply to the bulk lane only. The inference lane keeps the +// defaults. +const ( + bulkInitialWindow = 4 << 20 + bulkMaxWindow = 32 << 20 +) + +// sessionConfig returns the yamux configuration of a lane. It returns a new +// value for every call, because yamux keeps the pointer for the life of the +// session and two sessions must not share a struct. +// +// Replica-to-replica links carry the same kind of traffic as the bulk lane and +// use its windows. Both ends of a session must use the same configuration for +// it to take effect in both directions, because a window is announced by the +// side that receives. +func sessionConfig(lane Lane) *yamux.Config { + cfg := yamux.DefaultConfig() + // yamux writes its warnings to standard error. The callers report a dead + // session through the errors they get. + cfg.LogOutput = io.Discard + if lane == LaneBulk { + cfg.InitialStreamWindowSize = bulkInitialWindow + cfg.MaxStreamWindowSize = bulkMaxWindow + } + return cfg +} + +// Session is a multiplexed session on one websocket. +type Session struct { + s *yamux.Session +} + +// ServerSession starts the frontend side of a session on ws. The frontend owns +// the streams with even numbers and opens the streams. It does not accept any, +// because only the frontend asks and the worker answers. +func ServerSession(ws *websocket.Conn, lane Lane) (*Session, error) { + s, err := yamux.Server(websocketConn(ws), sessionConfig(lane), nil) + if err != nil { + return nil, fmt.Errorf("starting the server side of a tunnel session: %w", err) + } + return &Session{s: s}, nil +} + +// ClientSession starts the worker side of a session on ws. The worker accepts +// the streams that the frontend opens. +func ClientSession(ws *websocket.Conn, lane Lane) (*Session, error) { + s, err := yamux.Client(websocketConn(ws), sessionConfig(lane), nil) + if err != nil { + return nil, fmt.Errorf("starting the client side of a tunnel session: %w", err) + } + return &Session{s: s}, nil +} + +// OpenStream opens a stream to the other side. +func (s *Session) OpenStream(ctx context.Context) (net.Conn, error) { + st, err := s.s.OpenStream(ctx) + if err != nil { + return nil, err + } + return st, nil +} + +// AcceptStream waits for the other side to open a stream. It takes no context. +// Close the session to end the wait. +func (s *Session) AcceptStream() (net.Conn, error) { + st, err := s.s.AcceptStream() + if err != nil { + return nil, err + } + return st, nil +} + +// Close ends the session and every stream on it. +func (s *Session) Close() error { return s.s.Close() } + +// CloseChan is closed when the session has ended. +func (s *Session) CloseChan() <-chan struct{} { return s.s.CloseChan() } + +// IsClosed reports whether the session has ended. +func (s *Session) IsClosed() bool { return s.s.IsClosed() } diff --git a/core/services/tunnel/session_test.go b/core/services/tunnel/session_test.go new file mode 100644 index 000000000000..11cc3eb536a5 --- /dev/null +++ b/core/services/tunnel/session_test.go @@ -0,0 +1,154 @@ +package tunnel_test + +import ( + "context" + "io" + "time" + + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Buffers of the websocket", func() { + It("gives the upgrader and the dialer 64 KiB for reading and writing", func() { + up := tunnel.NewUpgrader() + Expect(up.ReadBufferSize).To(Equal(64 << 10)) + Expect(up.WriteBufferSize).To(Equal(64 << 10)) + + d := tunnel.NewDialer(3 * time.Second) + Expect(d.ReadBufferSize).To(Equal(64 << 10)) + Expect(d.WriteBufferSize).To(Equal(64 << 10)) + Expect(d.HandshakeTimeout).To(Equal(3 * time.Second)) + }) + + It("keeps the default origin check so that a browser of another origin is refused", func() { + Expect(tunnel.NewUpgrader().CheckOrigin).To(BeNil()) + }) +}) + +var _ = Describe("Lanes", func() { + It("reads the lane parameter, and reads an empty one as the inference lane", func() { + lane, err := tunnel.ParseLane("") + Expect(err).ToNot(HaveOccurred()) + Expect(lane).To(Equal(tunnel.LaneInference)) + + lane, err = tunnel.ParseLane("bulk") + Expect(err).ToNot(HaveOccurred()) + Expect(lane).To(Equal(tunnel.LaneBulk)) + + _, err = tunnel.ParseLane("express") + Expect(err).To(HaveOccurred()) + }) + + It("keeps the yamux defaults for the inference lane", func() { + cfg := tunnel.SessionConfigFor(tunnel.LaneInference) + Expect(cfg.InitialStreamWindowSize).To(BeEquivalentTo(256 << 10)) + Expect(cfg.MaxStreamWindowSize).To(BeEquivalentTo(16 << 20)) + }) + + It("uses windows of 4 MiB and 32 MiB for the bulk lane", func() { + cfg := tunnel.SessionConfigFor(tunnel.LaneBulk) + Expect(cfg.InitialStreamWindowSize).To(BeEquivalentTo(4 << 20)) + Expect(cfg.MaxStreamWindowSize).To(BeEquivalentTo(32 << 20)) + }) + + It("returns a new configuration for every call", func() { + a := tunnel.SessionConfigFor(tunnel.LaneBulk) + b := tunnel.SessionConfigFor(tunnel.LaneBulk) + Expect(a).ToNot(BeIdenticalTo(b)) + }) +}) + +var _ = Describe("Session", func() { + for _, lane := range []tunnel.Lane{tunnel.LaneInference, tunnel.LaneBulk} { + It("carries a stream from the frontend to the worker and back on the "+string(lane)+" lane", func() { + frontend, worker := sessionPair(lane) + + done := make(chan error, 1) + go func() { + st, err := worker.AcceptStream() + if err != nil { + done <- err + return + } + defer func() { _ = st.Close() }() + tag, target, err := tunnel.ReadStreamRequest(st) + if err != nil { + done <- err + return + } + if tag != tunnel.StreamTagGRPC || target != "127.0.0.1:50052" { + done <- io.ErrUnexpectedEOF + return + } + if err := tunnel.WriteStreamAccepted(st); err != nil { + done <- err + return + } + _, err = io.Copy(st, st) + done <- err + }() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + st, err := frontend.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(st, tunnel.StreamTagGRPC, "127.0.0.1:50052")).To(Succeed()) + Expect(tunnel.ReadStreamReply(st)).To(Succeed()) + _, err = st.Write([]byte("hello")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 5) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("hello")) + Expect(st.Close()).To(Succeed()) + Eventually(done).Should(Receive(BeNil())) + }) + } + + It("moves a payload larger than the buffers of the websocket without loss", func() { + frontend, worker := sessionPair(tunnel.LaneBulk) + payload := make([]byte, 5<<20+17) + for i := range payload { + payload[i] = byte(i * 31) + } + + echoed := make(chan error, 1) + go func() { + st, err := worker.AcceptStream() + if err != nil { + echoed <- err + return + } + _, err = io.Copy(st, st) + _ = st.Close() + echoed <- err + }() + + st, err := frontend.OpenStream(context.Background()) + Expect(err).ToNot(HaveOccurred()) + written := make(chan error, 1) + go func() { + _, werr := st.Write(payload) + written <- werr + }() + got := make([]byte, len(payload)) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal(payload)) + Eventually(written).Should(Receive(BeNil())) + Expect(st.Close()).To(Succeed()) + Eventually(echoed).Should(Receive()) + }) + + It("ends the other side when one side closes", func() { + frontend, worker := sessionPair(tunnel.LaneInference) + Expect(frontend.Close()).To(Succeed()) + Eventually(worker.CloseChan()).Should(BeClosed()) + Expect(worker.IsClosed()).To(BeTrue()) + _, err := worker.AcceptStream() + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/tunnel/tunnel_suite_test.go b/core/services/tunnel/tunnel_suite_test.go new file mode 100644 index 000000000000..97181fb3c958 --- /dev/null +++ b/core/services/tunnel/tunnel_suite_test.go @@ -0,0 +1,13 @@ +package tunnel_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestTunnel(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Tunnel test suite") +} diff --git a/core/services/tunnel/ws.go b/core/services/tunnel/ws.go new file mode 100644 index 000000000000..3173f08c5105 --- /dev/null +++ b/core/services/tunnel/ws.go @@ -0,0 +1,35 @@ +package tunnel + +import ( + "time" + + "github.com/gorilla/websocket" +) + +// BufferSize is the size of the read buffer and of the write buffer of every +// websocket of the tunnel. The default of gorilla is 4 KiB. With 4 KiB the +// direction that sends from the worker to the frontend was five times slower +// than the other direction on a loopback link, because the sending side masks +// every payload and does it in pieces of the size of the buffer. A buffer of +// 64 KiB matches the largest frame yamux sends. +const BufferSize = 64 << 10 + +// NewUpgrader returns the upgrader that every handler of the tunnel uses. It +// keeps the default origin check of gorilla, which refuses a browser from +// another origin and admits a client that sends no Origin header, as a worker +// does. +func NewUpgrader() *websocket.Upgrader { + return &websocket.Upgrader{ + ReadBufferSize: BufferSize, + WriteBufferSize: BufferSize, + } +} + +// NewDialer returns the dialer that every client of the tunnel uses. +func NewDialer(handshakeTimeout time.Duration) *websocket.Dialer { + return &websocket.Dialer{ + HandshakeTimeout: handshakeTimeout, + ReadBufferSize: BufferSize, + WriteBufferSize: BufferSize, + } +} diff --git a/core/services/tunnel/wsconn.go b/core/services/tunnel/wsconn.go new file mode 100644 index 000000000000..7a62cc5d976b --- /dev/null +++ b/core/services/tunnel/wsconn.go @@ -0,0 +1,134 @@ +package tunnel + +import ( + "errors" + "fmt" + "io" + "net" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// websocketConn adapts a gorilla websocket to the net.Conn that a yamux session +// drives. +// +// The two disagree about framing. A websocket delivers whole messages and yamux +// wants a stream of bytes. The adapter keeps the reader of the message that it +// is part-way through, so a Read with a buffer smaller than the message returns +// the first part now and the rest on the next call. This case is normal: yamux +// reads through a small buffered reader, and one stream write can put a much +// larger data frame on the wire in one message. +// +// The conn is safe for one reader and one writer at the same time, and for a +// third goroutine that sets deadlines. It is not a general net.Conn. +func websocketConn(ws *websocket.Conn) net.Conn { + return &wsConn{ws: ws} +} + +type wsConn struct { + ws *websocket.Conn + + // readMu guards frame, the partly read message. gorilla allows one reader, + // and this keeps to that even if a caller reads from two goroutines. + readMu sync.Mutex + frame io.Reader + + // writeMu keeps to the rule of gorilla that there is one writer. + writeMu sync.Mutex +} + +func (c *wsConn) Read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + + c.readMu.Lock() + defer c.readMu.Unlock() + + for { + if c.frame == nil { + messageType, r, err := c.ws.NextReader() + if err != nil { + return 0, translateReadErr(err) + } + // The link speaks binary messages only. If an unexpected text + // message were skipped, the yamux framing would go out of step + // without a sign, so it is an error. + if messageType != websocket.BinaryMessage { + return 0, fmt.Errorf("tunnel: received websocket message type %d, want binary", messageType) + } + c.frame = r + } + + n, err := c.frame.Read(p) + if errors.Is(err, io.EOF) { + // The end of one message is not the end of the stream. If io.EOF + // went up, the session would end at a message boundary. + c.frame = nil + err = nil + } + if n > 0 || err != nil { + return n, err + } + // A message of zero length gives nothing to return, and a Read that + // returns (0, nil) looks like a stalled stream to some callers. + } +} + +func (c *wsConn) Write(p []byte) (int, error) { + c.writeMu.Lock() + defer c.writeMu.Unlock() + + if err := c.ws.WriteMessage(websocket.BinaryMessage, p); err != nil { + return 0, err + } + return len(p), nil +} + +// Close drops the network connection without the websocket close handshake. +// yamux has sent its own go-away by then, and a close frame needs the write +// lock that a blocked send loop can still hold. +func (c *wsConn) Close() error { + return c.ws.Close() +} + +func (c *wsConn) LocalAddr() net.Addr { return c.ws.LocalAddr() } +func (c *wsConn) RemoteAddr() net.Addr { return c.ws.RemoteAddr() } + +func (c *wsConn) SetDeadline(t time.Time) error { + if err := c.SetReadDeadline(t); err != nil { + return err + } + return c.SetWriteDeadline(t) +} + +// SetReadDeadline takes no lock and must not take readMu. gorilla passes the +// deadline to the network connection, which allows this call from another +// goroutine, and readMu is held by the parked Read that this call must wake. +func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) } + +// SetWriteDeadline takes writeMu because gorilla keeps the write deadline in a +// plain field and applies it at the next flush. To set it during a write is a +// data race. +func (c *wsConn) SetWriteDeadline(t time.Time) error { + c.writeMu.Lock() + defer c.writeMu.Unlock() + + return c.ws.SetWriteDeadline(t) +} + +// translateReadErr maps a peer that hangs up in a clean way to io.EOF, which is +// how a yamux session knows a normal end. Other close codes and errors of the +// transport go up unchanged, so that the session reports a real failure. +func translateReadErr(err error) error { + if websocket.IsCloseError(err, + websocket.CloseNormalClosure, + websocket.CloseGoingAway, + websocket.CloseNoStatusReceived, + ) { + return io.EOF + } + return err +} diff --git a/go.mod b/go.mod index 342faf887ce0..4e2dcf35694f 100644 --- a/go.mod +++ b/go.mod @@ -33,6 +33,7 @@ require ( github.com/klauspost/cpuid/v2 v2.3.0 github.com/labstack/echo/v4 v4.15.1 github.com/libp2p/go-libp2p v0.48.0 + github.com/libp2p/go-yamux/v5 v5.1.0 github.com/lithammer/fuzzysearch v1.1.8 github.com/mholt/archiver/v3 v3.5.1 github.com/microcosm-cc/bluemonday v1.0.27 @@ -326,7 +327,6 @@ require ( github.com/jeandeaual/go-locale v0.0.0-20250612000132-0ef82f21eade // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/jsummers/gobmp v0.0.0-20230614200233-a9de23ed2e25 // indirect - github.com/libp2p/go-yamux/v5 v5.1.0 // indirect github.com/magiconair/properties v1.8.10 // indirect github.com/moby/docker-image-spec v1.3.1 // indirect github.com/moby/go-archive v0.2.0 // indirect From 9bc624536b8e88452fd98f9bc8a907d51770781f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:25:50 +0000 Subject: [PATCH 018/162] feat(tunnel): pass the end of a stream on when two streams are joined Splice used to close both streams as soon as one direction ended. A protocol that closes its sending side and then waits for the reply lost the reply that was still in flight. Now a direction that ends clean passes the end on with CloseWrite and the other direction goes on until it ends too. When the destination cannot half-close, or a direction fails or ends because a stream was closed, both streams are closed at once, as before. This is what wakes a copy that is parked in a read or in a write. A copy into a TCP connection wraps the error in a net.OpError. The classifier of the endings of yamux now removes that layer, so a stream that this side reset is not reported as a failure of the transport. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/splice.go | 197 ++++++++++++++++++++++++++++ core/services/tunnel/splice_test.go | 162 +++++++++++++++++++++++ 2 files changed, 359 insertions(+) create mode 100644 core/services/tunnel/splice.go create mode 100644 core/services/tunnel/splice_test.go diff --git a/core/services/tunnel/splice.go b/core/services/tunnel/splice.go new file mode 100644 index 000000000000..20945a41e560 --- /dev/null +++ b/core/services/tunnel/splice.go @@ -0,0 +1,197 @@ +package tunnel + +import ( + "errors" + "io" + "net" + + "github.com/libp2p/go-yamux/v5" +) + +// writeCloser is the half-close that *net.TCPConn and a yamux stream have. +type writeCloser interface { + CloseWrite() error +} + +// Splice joins two streams. It copies in both directions at the same time, +// because either side can speak first and gRPC has both directions live at +// once. Two copies in sequence would wait for a request on a stream whose other +// end waits for a response. +// +// When a direction reaches its end, Splice passes the end on with CloseWrite +// to the stream that the direction wrote to, and it goes on copying the other +// direction until that one ends as well. A protocol that closes its sending side +// and then waits for a reply, such as a request that ends in a half-close, so +// keeps its reply. If the stream cannot half-close, or the direction failed, +// Splice closes both streams at once. That is the only way to wake a copy that +// is parked in a read or in a write, and it is what the protocols that end a +// stream in both directions together (gRPC, HTTP/2) need when a peer is gone. +// +// Splice closes each stream exactly once and waits for both copies before it +// returns, so no copy touches a stream after the return. The wait assumes that +// Close wakes a copy that is parked in Read or Write. A net.Conn and a yamux +// stream do that. Check this before using another type. +// +// The result is the first failure of a direction. An ending that means "someone +// closed" is not a failure. A result therefore means "one direction failed" and +// never "only this failed". +func Splice(a, b io.ReadWriteCloser) error { + results := make(chan copyResult, 2) + go func() { results <- copyStream(b, a) }() + go func() { results <- copyStream(a, b) }() + + first := <-results + if first.halfClosed { + // The end was passed on. The other direction ends when the peer closes + // its side, or it fails when the session dies, and either is its own + // result and not an echo of a close by Splice. + second := <-results + closeErrA := a.Close() + closeErrB := b.Close() + if second.err != nil { + return second.err + } + return firstCloseFailure(closeErrA, closeErrB) + } + + // Each stream is closed here and nowhere else. A second Close of a yamux + // stream that was reset returns the error that killed it, and Splice could + // not tell that from a new failure. + closeErrA := a.Close() + closeErrB := b.Close() + + // The other direction ends because of the Close above. Its result is an + // echo of that Close, so it is dropped, and the failure of the first + // direction is the one to report. It can also be a real failure that lost + // the race, which is why a result means "one direction failed" and never + // "only this failed". + <-results + + if first.err != nil { + return first.err + } + return firstCloseFailure(closeErrA, closeErrB) +} + +// firstCloseFailure returns the first Close that failed on a stream that was +// healthy. A close of a stream that is already dead is not a failure. +func firstCloseFailure(errA, errB error) error { + if err := normalizeStreamErr(errA); err != nil { + return err + } + return normalizeStreamErr(errB) +} + +// copyResult is the end of one direction. halfClosed is true when the end was +// passed on to the destination with CloseWrite. +type copyResult struct { + err error + halfClosed bool +} + +// copyStream moves one direction and reports only failures of the transport. +func copyStream(dst io.Writer, src io.Reader) copyResult { + _, err := io.Copy(dst, src) + if err != nil { + // A read that ends because a stream was closed or reset is not the end + // of the data. The peer may never close its side, so Splice must close + // both streams and not wait for it. + return copyResult{err: normalizeStreamErr(err)} + } + cw, ok := dst.(writeCloser) + if !ok { + return copyResult{} + } + if err := normalizeStreamErr(cw.CloseWrite()); err != nil { + return copyResult{err: err} + } + return copyResult{halfClosed: true} +} + +// normalizeStreamErr drops the endings that mean that the conversation is over +// and not broken: net.ErrClosed, which a socket reports after it or its peer +// closed, and io.ErrClosedPipe, which is the same on an in-memory pipe. +// +// io.EOF is not in the list. io.Copy consumes a clean end of a read and reports +// nil, so an io.EOF that arrives here came from a Write or a Close that failed, +// and it means that the peer is gone. yamux gives exactly that when a Write +// races the shutdown of its session. +// +// A reset of a socket (ECONNRESET, EPIPE) is not in the list either, so the +// same event, a peer that aborts in the middle of a stream, is an error on a +// raw socket and, through muxVerdict, on a yamux stream. +// +// The verdict of the multiplexer comes first. A session that dies gives every +// live stream its own cause, and that cause is often a closed-socket error, so +// the generic endings would report a peer that vanished as a clean end. +func normalizeStreamErr(err error) error { + if err == nil { + return nil + } + if recognised, report := muxVerdict(err); recognised { + if report { + return err + } + return nil + } + if errors.Is(err, net.ErrClosed) || errors.Is(err, io.ErrClosedPipe) { + return nil + } + return err +} + +// normalGoAwayCode is the code of a go-away that reports no error. yamux does +// not export the constant, so it is read from a sentinel that has it. +var normalGoAwayCode = yamux.ErrRemoteGoAway.ErrorCode + +// muxVerdict classifies an ending that comes from yamux. recognised says that +// the error is from the multiplexer. report says that the ending was forced on +// this stream and was not asked for by this side. The policy turns on one bit, +// Remote, and it is in one function so that a test can show each read of it. +// +// - A reset of the stream by the peer is reported. For a request that is +// relayed it means that the response is cut. A caller decides how loud that +// is: a client that cancels causes one for each cancel. +// - A go-away from the peer is reported. It is graceful for the session, but +// every stream on the session was in the middle of a request. +// - The same endings caused by this side are silent. ErrSessionShutdown is +// what this process gets when it closes its own session. +// +// A reset that arrives on a frame type other than a window update gives the bare +// ErrStreamReset sentinel, which is treated as the teardown of this side. Every +// reset that go-yamux sends uses a window update, so only a foreign +// implementation of the protocol can reach that case. +func muxVerdict(err error) (recognised, report bool) { + var goAway *yamux.GoAwayError + if errors.As(err, &goAway) { + return true, goAway.Remote || goAway.ErrorCode != normalGoAwayCode + } + var streamErr *yamux.StreamError + if errors.As(err, &streamErr) { + return true, streamErr.Remote + } + // These sentinels are compared by identity and not with errors.Is. They are + // the endings that Splice causes itself: a Write on a stream that was closed + // gives ErrStreamClosed, and a Read that was parked when CloseRead ran gives + // the bare ErrStreamReset. + // + // A copy into a *net.TCPConn wraps the error in a net.OpError, so the layer + // is removed before the comparison. + for { + opErr, ok := err.(*net.OpError) + if !ok { + break + } + err = opErr.Err + } + if err == yamux.ErrStreamClosed || err == yamux.ErrStreamReset { + return true, false + } + // The same sentinel wrapped is another thing: Session.close gives every + // stream that it kills ErrStreamReset around the cause. The session died + // under a live stream. Identity above is what tells the two apart. + if errors.Is(err, yamux.ErrStreamReset) { + return true, true + } + return false, false +} diff --git a/core/services/tunnel/splice_test.go b/core/services/tunnel/splice_test.go new file mode 100644 index 000000000000..f14f3166f0bf --- /dev/null +++ b/core/services/tunnel/splice_test.go @@ -0,0 +1,162 @@ +package tunnel_test + +import ( + "context" + "errors" + "io" + "net" + "time" + + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// streamAndTCP returns a stream of a session (the stream end that Splice gets +// on the worker) with the other end of that stream, and a TCP connection with +// the other end of that connection. +func streamAndTCP() (far net.Conn, spliceStream net.Conn, spliceTCP, backend *net.TCPConn) { + GinkgoHelper() + frontend, worker := sessionPair(tunnel.LaneInference) + type accepted struct { + c net.Conn + err error + } + ch := make(chan accepted, 1) + go func() { + c, err := worker.AcceptStream() + ch <- accepted{c, err} + }() + far, err := frontend.OpenStream(context.Background()) + Expect(err).ToNot(HaveOccurred()) + // A stream is announced to the other side with its first bytes. + Expect(tunnel.WriteStreamAccepted(far)).To(Succeed()) + var got accepted + Eventually(ch).Should(Receive(&got)) + Expect(got.err).ToNot(HaveOccurred()) + Expect(tunnel.ReadStreamReply(got.c)).To(Succeed()) + + spliceTCP, backend = tcpPair() + return far, got.c, spliceTCP, backend +} + +var _ = Describe("Splice", func() { + It("carries both directions at the same time", func() { + far, stream, local, backend := streamAndTCP() + done := make(chan error, 1) + go func() { done <- tunnel.Splice(stream, local) }() + + // The backend speaks first, as a server of gRPC does after the preface. + _, err := backend.Write([]byte("settings")) + Expect(err).ToNot(HaveOccurred()) + _, err = far.Write([]byte("preface")) + Expect(err).ToNot(HaveOccurred()) + + got := make([]byte, 8) + _, err = io.ReadFull(far, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("settings")) + got = make([]byte, 7) + _, err = io.ReadFull(backend, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("preface")) + + Expect(far.Close()).To(Succeed()) + _ = backend.Close() + Eventually(done, 5*time.Second).Should(Receive(BeNil())) + }) + + It("keeps the reply of a request that ends in a half-close", func() { + far, stream, local, backend := streamAndTCP() + done := make(chan error, 1) + go func() { done <- tunnel.Splice(stream, local) }() + + // The backend reads the request up to its end and only then answers. + go func() { + defer GinkgoRecover() + req, err := io.ReadAll(backend) + Expect(err).ToNot(HaveOccurred()) + _, err = backend.Write(append([]byte("reply to "), req...)) + Expect(err).ToNot(HaveOccurred()) + Expect(backend.CloseWrite()).To(Succeed()) + }() + + _, err := far.Write([]byte("request")) + Expect(err).ToNot(HaveOccurred()) + cw, ok := far.(interface{ CloseWrite() error }) + Expect(ok).To(BeTrue()) + Expect(cw.CloseWrite()).To(Succeed()) + + reply, err := io.ReadAll(far) + Expect(err).ToNot(HaveOccurred()) + Expect(string(reply)).To(Equal("reply to request")) + + Expect(far.Close()).To(Succeed()) + Eventually(done, 5*time.Second).Should(Receive(BeNil())) + }) + + It("passes the half-close of the backend to the stream and keeps the request direction open", func() { + far, stream, local, backend := streamAndTCP() + done := make(chan error, 1) + go func() { done <- tunnel.Splice(stream, local) }() + + Expect(backend.CloseWrite()).To(Succeed()) + // The far side reads the end of the stream, and it can still write. + _, err := io.ReadAll(far) + Expect(err).ToNot(HaveOccurred()) + _, err = far.Write([]byte("late")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 4) + _, err = io.ReadFull(backend, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("late")) + + Expect(far.Close()).To(Succeed()) + Eventually(done, 5*time.Second).Should(Receive(BeNil())) + }) + + It("closes both streams when one end cannot half-close", func() { + a1, a2 := net.Pipe() + b1, b2 := net.Pipe() + done := make(chan error, 1) + go func() { done <- tunnel.Splice(a2, b1) }() + + Expect(a1.Close()).To(Succeed()) + Eventually(done, 5*time.Second).Should(Receive(BeNil())) + // The other end was closed by Splice, so its peer sees the end. + _, err := b2.Read(make([]byte, 1)) + Expect(err).To(MatchError(io.EOF)) + }) + + It("closes both streams at once when a direction fails", func() { + far, stream, local, backend := streamAndTCP() + done := make(chan error, 1) + go func() { done <- tunnel.Splice(stream, local) }() + + // Closing the whole session under a live stream is a failure of the + // transport and not an end that was asked for. + Expect(stream.Close()).To(Succeed()) + Eventually(done, 5*time.Second).Should(Receive(BeNil())) + _, err := backend.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred()) + _ = far.Close() + }) + + It("reports a stream that the peer reset in the middle of a request", func() { + a, b := net.Pipe() + boom := errors.New("link broke") + done := make(chan error, 1) + go func() { done <- tunnel.Splice(&failingConn{Conn: a, err: boom}, b) }() + _, _ = b.Write([]byte("x")) + Eventually(done, 5*time.Second).Should(Receive(MatchError(boom))) + }) +}) + +// failingConn fails every read after the first byte. +type failingConn struct { + net.Conn + err error +} + +func (f *failingConn) Read(p []byte) (int, error) { return 0, f.err } From c5b00abada88746bf46e89a7a54cb6d74ad1246d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:25:50 +0000 Subject: [PATCH 019/162] test(tunnel): measure the delay of a small call during a large transfer A proxy in the spec works as one link of 1 Gbit/s with 10 ms of delay in each direction and one queue that every connection shares. A transfer of 256 MiB runs from the worker to the frontend while a small call goes every 5 ms in the other direction, and its reply has to wait in the queue of the loaded direction. The spec runs three set-ups: a probe on its own TCP connection, a probe on the session of the transfer, and a probe on the inference session with the transfer on the bulk session. It requires the bulk lane to stay within 10 ms of the TCP baseline at the 99th percentile, and the shared session to be worse than the bulk lane. A build with the race detector moves 128 MiB. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/latency_test.go | 483 ++++++++++++++++++++++++++ core/services/tunnel/race_off_test.go | 8 + core/services/tunnel/race_on_test.go | 5 + 3 files changed, 496 insertions(+) create mode 100644 core/services/tunnel/latency_test.go create mode 100644 core/services/tunnel/race_off_test.go create mode 100644 core/services/tunnel/race_on_test.go diff --git a/core/services/tunnel/latency_test.go b/core/services/tunnel/latency_test.go new file mode 100644 index 000000000000..c1c6002f4808 --- /dev/null +++ b/core/services/tunnel/latency_test.go @@ -0,0 +1,483 @@ +package tunnel_test + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "slices" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// slowLink is a TCP proxy that works as one bottleneck link for every +// connection that goes through it. It has a rate and a delay in each +// direction, and one queue of 256 KiB for each direction that all connections +// share. A real link has that shape, and the effect under test depends on it: a +// probe that has its own connection still waits in the queue behind the bytes +// of a transfer on another connection. What a probe must not wait for are the +// bytes that sit in the buffers of its own connection ahead of it. +type slowLink struct { + lis net.Listener + target string + rate float64 + oneWay time.Duration + up *linkQueue + down *linkQueue + quit chan struct{} + wg sync.WaitGroup + + mu sync.Mutex + conns []net.Conn +} + +type linkItem struct { + data []byte + dst net.Conn + at time.Time +} + +// linkQueue is one direction of the link. +type linkQueue struct { + in chan linkItem + delivery chan linkItem +} + +func newSlowLink(target string, bytesPerSecond float64, oneWay time.Duration) *slowLink { + GinkgoHelper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + l := &slowLink{ + lis: lis, target: target, rate: bytesPerSecond, oneWay: oneWay, + up: &linkQueue{in: make(chan linkItem, 4), delivery: make(chan linkItem, 4096)}, + down: &linkQueue{in: make(chan linkItem, 4), delivery: make(chan linkItem, 4096)}, + quit: make(chan struct{}), + } + for _, q := range []*linkQueue{l.up, l.down} { + l.wg.Add(2) + go l.transmit(q) + go l.deliver(q) + } + go l.accept() + DeferCleanup(l.close) + return l +} + +func (l *slowLink) addr() string { return l.lis.Addr().String() } + +func (l *slowLink) close() { + select { + case <-l.quit: + return + default: + } + close(l.quit) + _ = l.lis.Close() + l.mu.Lock() + for _, c := range l.conns { + _ = c.Close() + } + l.mu.Unlock() + l.wg.Wait() +} + +func (l *slowLink) accept() { + for { + c, err := l.lis.Accept() + if err != nil { + return + } + s, err := net.Dial("tcp", l.target) + if err != nil { + _ = c.Close() + continue + } + l.mu.Lock() + l.conns = append(l.conns, c, s) + l.mu.Unlock() + go l.read(c, s, l.up) + go l.read(s, c, l.down) + } +} + +// read moves what arrives on src into the queue of a direction. A full queue +// stops the read, so the socket fills and the sender slows down, as it would +// on a link with a queue of limited size. +func (l *slowLink) read(src, dst net.Conn, q *linkQueue) { + defer func() { _ = dst.Close() }() + for { + buf := make([]byte, 64<<10) + n, err := src.Read(buf) + if n > 0 { + select { + case q.in <- linkItem{data: buf[:n], dst: dst}: + case <-l.quit: + return + } + } + if err != nil { + return + } + } +} + +// transmit takes an item out of the queue at the rate of the link. +func (l *slowLink) transmit(q *linkQueue) { + defer l.wg.Done() + next := time.Now() + for { + var item linkItem + select { + case item = <-q.in: + case <-l.quit: + return + } + // A late wake-up must not slow the link: the schedule catches up + // instead. Only an idle link starts a new schedule. + if now := time.Now(); now.Sub(next) > 5*time.Millisecond { + next = now + } + next = next.Add(time.Duration(float64(len(item.data)) / l.rate * float64(time.Second))) + if d := time.Until(next); d > 500*time.Microsecond { + time.Sleep(d) + } + item.at = time.Now().Add(l.oneWay) + select { + case q.delivery <- item: + case <-l.quit: + return + } + } +} + +// deliver writes an item to its connection after the delay of the link. One +// goroutine for each direction keeps the order of each connection. +func (l *slowLink) deliver(q *linkQueue) { + defer l.wg.Done() + for { + var item linkItem + select { + case item = <-q.delivery: + case <-l.quit: + return + } + time.Sleep(time.Until(item.at)) + _, _ = item.dst.Write(item.data) + } +} + +// probeStats are the delays of the probes in milliseconds. +type probeStats struct { + count int + p50, p99, max float64 +} + +func summarize(samples []time.Duration) probeStats { + if len(samples) == 0 { + return probeStats{} + } + slices.Sort(samples) + ms := func(d time.Duration) float64 { return float64(d) / float64(time.Millisecond) } + return probeStats{ + count: len(samples), + p50: ms(samples[len(samples)/2]), + p99: ms(samples[min(len(samples)-1, len(samples)*99/100)]), + max: ms(samples[len(samples)-1]), + } +} + +// probeWhile calls probe every few milliseconds until stop is closed. +func probeWhile(stop <-chan struct{}, probe func() (time.Duration, error)) []time.Duration { + var samples []time.Duration + for { + select { + case <-stop: + return samples + case <-time.After(5 * time.Millisecond): + } + d, err := probe() + if err != nil { + return samples + } + samples = append(samples, d) + } +} + +// echoTCP answers every byte that it reads with the same byte. +func echoTCP(c net.Conn) { + defer func() { _ = c.Close() }() + _, _ = io.Copy(c, c) +} + +// sinkTCP reads and drops everything. +func sinkTCP(c net.Conn) { + defer func() { _ = c.Close() }() + _, _ = io.Copy(io.Discard, c) +} + +// tcpService starts a server that runs handle for every connection. +func tcpService(handle func(net.Conn)) string { + GinkgoHelper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(lis.Close) + go func() { + for { + c, err := lis.Accept() + if err != nil { + return + } + go handle(c) + } + }() + return lis.Addr().String() +} + +func pushBytes(w io.Writer, size int) error { + chunk := make([]byte, 256<<10) + for sent := 0; sent < size; sent += len(chunk) { + if _, err := w.Write(chunk); err != nil { + return err + } + } + return nil +} + +// tcpProbe returns a probe that sends 8 bytes on one connection and waits for +// the echo. +func tcpProbe(c net.Conn) func() (time.Duration, error) { + req, rep := make([]byte, 8), make([]byte, 8) + return func() (time.Duration, error) { + start := time.Now() + if _, err := c.Write(req); err != nil { + return 0, err + } + if _, err := io.ReadFull(c, rep); err != nil { + return 0, err + } + return time.Since(start), nil + } +} + +// lanedFrontend is a frontend that serves websockets and keeps the sessions +// that it gets, by lane, as the connect endpoint does. +type lanedFrontend struct { + addr string + inference chan *tunnel.Session + bulk chan *tunnel.Session +} + +func startLanedFrontend() *lanedFrontend { + GinkgoHelper() + f := &lanedFrontend{inference: make(chan *tunnel.Session, 2), bulk: make(chan *tunnel.Session, 2)} + up := tunnel.NewUpgrader() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + lane, err := tunnel.ParseLane(r.URL.Query().Get("lane")) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + ws, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + s, err := tunnel.ServerSession(ws, lane) + if err != nil { + _ = ws.Close() + return + } + if lane == tunnel.LaneBulk { + f.bulk <- s + } else { + f.inference <- s + } + })) + DeferCleanup(srv.Close) + f.addr = srv.Listener.Addr().String() + return f +} + +// dialThrough opens the worker side of a session on lane, through the link. +func dialThrough(linkAddr string, lane tunnel.Lane) *tunnel.Session { + GinkgoHelper() + ws, _, err := tunnel.NewDialer(10*time.Second).Dial("ws://"+linkAddr+"/?lane="+string(lane), nil) + Expect(err).ToNot(HaveOccurred()) + worker, err := tunnel.ClientSession(ws, lane) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = worker.Close() }) + return worker +} + +// serveStreams runs handle for every stream that the other side opens. +func serveStreams(s *tunnel.Session, handle func(net.Conn)) { + go func() { + for { + st, err := s.AcceptStream() + if err != nil { + return + } + go handle(st) + } + }() +} + +// streamProbe returns a probe that opens a stream on a session, sends 8 bytes +// and waits for the echo. +func streamProbe(s *tunnel.Session) func() (time.Duration, error) { + req, rep := make([]byte, 8), make([]byte, 8) + return func() (time.Duration, error) { + start := time.Now() + st, err := s.OpenStream(context.Background()) + if err != nil { + return 0, err + } + defer func() { _ = st.Close() }() + if _, err := st.Write(req); err != nil { + return 0, err + } + if _, err := io.ReadFull(st, rep); err != nil { + return 0, err + } + return time.Since(start), nil + } +} + +// The transfer goes from the worker to the frontend, which is the direction +// that the benchmark of the buffers found to be the slow one. The probe is a +// call from the frontend to the worker: its request travels in the direction +// with no load, and its reply travels in the direction of the transfer, where +// it has to wait behind the bytes that are already queued. +var _ = Describe("Delay of a small call during a large transfer", Label("benchmark"), Ordered, func() { + const ( + linkRate = 125e6 // 1 Gbit/s + linkOneWay = 10 * time.Millisecond + + directName = "direct: transfer and probe on two TCP connections" + sharedName = "shared session: transfer and probe on one session" + bulkName = "bulk lane: transfer on the bulk session, probe on the inference session" + ) + results := map[string]probeStats{} + + // measure runs the transfer while the probes run, and records the stats of + // the probes. + measure := func(name string, transfer func() error, probe func() (time.Duration, error)) { + GinkgoHelper() + stop := make(chan struct{}) + sampled := make(chan []time.Duration, 1) + go func() { sampled <- probeWhile(stop, probe) }() + start := time.Now() + Expect(transfer()).To(Succeed()) + elapsed := time.Since(start) + close(stop) + var samples []time.Duration + Eventually(sampled, 30*time.Second).Should(Receive(&samples)) + Expect(len(samples)).To(BeNumerically(">=", 15), "too few probes completed during the transfer") + // The first probes run while the transfer fills the link, so the first + // tenth does not count. + stats := summarize(samples[len(samples)/10:]) + results[name] = stats + AddReportEntry(name, map[string]any{ + "transfer MiB": transferSize >> 20, + "transfer s": elapsed.Seconds(), + "transfer MiB/s": float64(transferSize>>20) / elapsed.Seconds(), + "probes": stats.count, + "probe p50 ms": stats.p50, + "probe p99 ms": stats.p99, + "probe max ms": stats.max, + }) + } + + It("measures the baseline: the probe and the transfer on two TCP connections", func() { + // One service behind the link. The first byte of a connection says what + // the connection is for. + front := tcpService(func(c net.Conn) { + kind := make([]byte, 1) + if _, err := io.ReadFull(c, kind); err != nil { + _ = c.Close() + return + } + if kind[0] == 'e' { + echoTCP(c) + } else { + sinkTCP(c) + } + }) + link := newSlowLink(front, linkRate, linkOneWay) + + dial := func(kind byte) net.Conn { + c, err := net.Dial("tcp", link.addr()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(c.Close) + _, err = c.Write([]byte{kind}) + Expect(err).ToNot(HaveOccurred()) + return c + } + bulk, probe := dial('s'), dial('e') + + measure(directName, + func() error { return pushBytes(bulk, transferSize) }, + tcpProbe(probe)) + }) + + It("measures the probe on the same session as the transfer", func() { + f := startLanedFrontend() + link := newSlowLink(f.addr, linkRate, linkOneWay) + worker := dialThrough(link.addr(), tunnel.LaneInference) + var frontend *tunnel.Session + Eventually(f.inference).Should(Receive(&frontend)) + DeferCleanup(func() { _ = frontend.Close() }) + serveStreams(worker, echoTCP) + serveStreams(frontend, sinkTCP) + + measure(sharedName, + func() error { + st, err := worker.OpenStream(context.Background()) + if err != nil { + return err + } + defer func() { _ = st.Close() }() + return pushBytes(st, transferSize) + }, + streamProbe(frontend)) + }) + + It("measures the probe on the inference lane with the transfer on the bulk lane", func() { + f := startLanedFrontend() + // Both sessions go through the same link, as the two websockets of one + // worker do. + link := newSlowLink(f.addr, linkRate, linkOneWay) + inferenceWorker := dialThrough(link.addr(), tunnel.LaneInference) + bulkWorker := dialThrough(link.addr(), tunnel.LaneBulk) + var inferenceFrontend, bulkFrontend *tunnel.Session + Eventually(f.inference).Should(Receive(&inferenceFrontend)) + Eventually(f.bulk).Should(Receive(&bulkFrontend)) + DeferCleanup(func() { _ = inferenceFrontend.Close(); _ = bulkFrontend.Close() }) + serveStreams(inferenceWorker, echoTCP) + serveStreams(bulkFrontend, sinkTCP) + + measure(bulkName, + func() error { + st, err := bulkWorker.OpenStream(context.Background()) + if err != nil { + return err + } + defer func() { _ = st.Close() }() + return pushBytes(st, transferSize) + }, + streamProbe(inferenceFrontend)) + }) + + It("keeps the probe on the bulk lane within 10 ms of the direct baseline at the 99th percentile", func() { + direct, bulk, shared := results[directName], results[bulkName], results[sharedName] + Expect(bulk.p99).To(BeNumerically("<=", direct.p99+10), + "bulk lane p99 %.1f ms, direct p99 %.1f ms", bulk.p99, direct.p99) + Expect(shared.p99).To(BeNumerically(">", bulk.p99), + "a shared session must delay the probe more than the bulk lane: shared p99 %.1f ms, bulk p99 %.1f ms", shared.p99, bulk.p99) + }) +}) diff --git a/core/services/tunnel/race_off_test.go b/core/services/tunnel/race_off_test.go new file mode 100644 index 000000000000..77e990ad725a --- /dev/null +++ b/core/services/tunnel/race_off_test.go @@ -0,0 +1,8 @@ +//go:build !race + +package tunnel_test + +// transferSize is the size of the transfer in the latency spec. The race +// detector slows the copy so much that the transfer would take most of a +// minute, so a build with it moves less (see race_on_test.go). +const transferSize = 256 << 20 diff --git a/core/services/tunnel/race_on_test.go b/core/services/tunnel/race_on_test.go new file mode 100644 index 000000000000..961a82a0353b --- /dev/null +++ b/core/services/tunnel/race_on_test.go @@ -0,0 +1,5 @@ +//go:build race + +package tunnel_test + +const transferSize = 128 << 20 From a27b9111c72c8ea3b98e3cae815c9f0af109057b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:38:53 +0000 Subject: [PATCH 020/162] feat(tunnel): hold the tunnels of a replica and tie each to a claim The registry keeps the sessions that workers opened to this replica and keeps the node_connections table in agreement with them. Attach writes the claim first and stores the session second, so a failed claim leaves the replica as it was, and the claim and its record are one step for each node. Detach matches the token of the attachment by equality, so a session that was replaced cannot remove the one that replaced it. Open returns ErrNotOwner only when the node is not held here. A session that ended under a held entry, and a caller whose budget ran out, are reported as themselves. A node has two lanes. The inference lane owns the claim. The bulk lane is stored under the claim of the inference lane, is refused with ErrNotOwner on a replica that does not hold it, and is closed with it. Open on the bulk lane uses the inference lane when the node has no bulk session, so a transfer on an older worker or during a reconnect works. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/registry.go | 501 ++++++++++++++++++++++++++ core/services/tunnel/registry_test.go | 414 +++++++++++++++++++++ 2 files changed, 915 insertions(+) create mode 100644 core/services/tunnel/registry.go create mode 100644 core/services/tunnel/registry_test.go diff --git a/core/services/tunnel/registry.go b/core/services/tunnel/registry.go new file mode 100644 index 000000000000..6450cbbae323 --- /dev/null +++ b/core/services/tunnel/registry.go @@ -0,0 +1,501 @@ +package tunnel + +import ( + "context" + "errors" + "fmt" + "net" + "sort" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/xlog" +) + +// ConnectPath is the route that a worker dials to open its tunnel. The HTTP +// layer registers the handler on it and the worker builds its URL from it. +const ConnectPath = "/api/cluster/connect" + +// ErrNotOwner means that this replica does not hold the tunnel of a node. +// +// It is a routing fact and nothing else: another replica can hold that worker, +// and a caller that gets this error asks the owner that the database names. It +// must never come from anything that only failed. A database error, a broken +// socket and a session that ended under a held entry are all reported as +// themselves, because "not held here" tells a dialer to look elsewhere for a +// worker that is right here, and absence must never stand in for a failure. +var ErrNotOwner = errors.New("tunnel: this replica does not hold the tunnel for that node") + +// releaseTimeout bounds the release in Detach. Detach runs in the goroutine that +// has just seen a session end, and that goroutine must not wait for a database +// that went away with the session. +const releaseTimeout = 5 * time.Second + +// Registry holds the tunnels that this replica has accepted, and keeps the +// node_connections table in agreement with what it holds. +// +// The table says which replica owns a worker. This registry says which socket +// that ownership means. Attach writes the two in one order, always. +// +// A node has up to two sessions: the inference lane and the bulk lane. They +// share one claim and one epoch. The claim belongs to the inference lane, which +// is the session that the worker must have. The bulk lane is an addition to it +// and is held only on the replica that holds the inference lane. +type Registry struct { + reg *cluster.Registry + selfID string + + mu sync.Mutex + tunnels map[string]*heldTunnel + // claiming holds one gate for each node that has a claim in progress. It + // makes "claim and then record the claim" one step for each node. + claiming map[string]chan struct{} + // bulkTokens numbers the attachments of the bulk lane. A bulk lane has no + // claim of its own, so its attachment needs its own identity for Detach. + bulkTokens int64 +} + +// heldTunnel is the tunnel of one node. +// +// token identifies one attachment of the inference lane for its whole life, and +// it is the only value that Detach compares. claim is the epoch of the row that +// this replica holds for the node, and Release needs it because the fence +// matches the row exactly. The two are the same number until this replica is +// swept and claims again, and they stay separate because they answer different +// questions. Neither is compared for order: Claim promises uniqueness, not +// growth. +type heldTunnel struct { + inference *Session + token int64 + claim int64 + + bulk *Session + bulkToken int64 +} + +// NewRegistry returns a registry that claims tunnels as selfID. The ID must be +// the one that this replica registers in the instances table, because Owner +// joins a claim against that table to decide if the owner is alive. +func NewRegistry(reg *cluster.Registry, selfID string) *Registry { + return &Registry{ + reg: reg, + selfID: selfID, + tunnels: map[string]*heldTunnel{}, + claiming: map[string]chan struct{}{}, + } +} + +// enterClaim takes the gate of a node, so that two claims of one node are never +// in progress together. leaveClaim gives it back. +// +// A claim and the record of that claim are two steps, and the database moves +// between them. Two Attach calls for one node both claim, and PostgreSQL puts +// the two writes in order, but nothing orders the two map writes against the +// two commits. The entry that ends up in the map could carry the epoch of the +// claim that did not win the row. Its Detach would then release an epoch that +// the row does not hold, the release would match nothing, and the row would +// outlive the socket. Nothing sweeps that row, because the replica named on it +// is alive. Owner would go on naming this replica as the owner of a tunnel that +// it does not hold. +// +// The gate is per node and not for the whole registry, so a slow claim of one +// worker does not hold up Open for the others. Detach does not use the gate. It +// takes no context and must not wait for a database call. It does not need the +// gate, because it changes no epoch. +func (t *Registry) enterClaim(ctx context.Context, nodeID string) error { + for { + t.mu.Lock() + gate, busy := t.claiming[nodeID] + if !busy { + t.claiming[nodeID] = make(chan struct{}) + t.mu.Unlock() + return nil + } + t.mu.Unlock() + + // The loop checks again after a wake-up. One close wakes several + // waiters and only one of them may go on. + select { + case <-gate: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// leaveClaim gives back the gate. The channel is closed, so every waiter wakes. +func (t *Registry) leaveClaim(nodeID string) { + t.mu.Lock() + gate := t.claiming[nodeID] + delete(t.claiming, nodeID) + t.mu.Unlock() + close(gate) +} + +// Attach records this replica as the owner of the tunnel of nodeID and stores +// the session. The result is the token that the caller must give to Detach. +// +// For the inference lane the claim is written before the session is stored. +// A claimant that installs itself first and finds later that it cannot claim +// has published, for a while, a tunnel that no row records: Held would name it, +// and another replica that asks Owner would be told that the worker is not +// connected anywhere. If the claim comes first, a claim that fails leaves this +// replica as it was. +// +// A worker that dials again to the same replica replaces its earlier +// attachment, and the registry closes the session that it replaced. Nothing +// else can close it: the goroutine that accepted it is parked on a session that +// is not broken, only replaced. A bulk session of the replaced attachment is +// closed too. The caller keeps the session that it passed in. +// +// The bulk lane has no claim. It is stored under the tunnel that this replica +// holds, and without one the result is ErrNotOwner. The worker is then told to +// try another replica, and its next dial can land on the owner. +func (t *Registry) Attach(ctx context.Context, nodeID string, lane Lane, sess *Session) (int64, error) { + if sess == nil { + return 0, fmt.Errorf("attaching tunnel for node %q: no session", nodeID) + } + switch lane { + case LaneInference: + return t.attachInference(ctx, nodeID, sess) + case LaneBulk: + return t.attachBulk(nodeID, sess) + default: + return 0, fmt.Errorf("attaching tunnel for node %q: unknown lane %q", nodeID, lane) + } +} + +func (t *Registry) attachInference(ctx context.Context, nodeID string, sess *Session) (int64, error) { + if err := t.enterClaim(ctx, nodeID); err != nil { + return 0, fmt.Errorf("attaching tunnel for node %q: %w", nodeID, err) + } + + // The part under the gate is a closure, so that the gate is released with a + // defer and the close of the old session below still happens outside the + // gate. A panic under Claim, which does database work, would otherwise leave + // the gate of this node closed for the life of the process. + var previous *heldTunnel + epoch, err := func() (int64, error) { + defer t.leaveClaim(nodeID) + + epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) + if err != nil { + return 0, err + } + + t.mu.Lock() + previous = t.tunnels[nodeID] + t.tunnels[nodeID] = &heldTunnel{inference: sess, token: epoch, claim: epoch} + t.mu.Unlock() + return epoch, nil + }() + if err != nil { + return 0, err + } + + // Closed after the gate is released. Yamux closes the connection and waits + // for its send and receive loops, and the send loop can be in a write that + // only the write timeout bounds. That wait is on other goroutines and must + // not stand between a worker that dials again and its claim. This is safe + // because the old session can no longer be reached from the map. + if previous != nil { + if previous.inference != sess { + xlog.Debug("worker dialled this replica again, dropping its previous tunnel", "node", nodeID) + _ = previous.inference.Close() + } + if previous.bulk != nil { + _ = previous.bulk.Close() + } + } + return epoch, nil +} + +func (t *Registry) attachBulk(nodeID string, sess *Session) (int64, error) { + t.mu.Lock() + held, ok := t.tunnels[nodeID] + if !ok || held.inference.IsClosed() { + t.mu.Unlock() + return 0, fmt.Errorf("attaching the bulk lane for node %q: %w", nodeID, ErrNotOwner) + } + previous := held.bulk + t.bulkTokens++ + held.bulk = sess + held.bulkToken = t.bulkTokens + token := held.bulkToken + t.mu.Unlock() + + if previous != nil && previous != sess { + xlog.Debug("worker dialled the bulk lane again, dropping its previous session", "node", nodeID) + _ = previous.Close() + } + return token, nil +} + +// Detach drops the attachment that token identifies. For the inference lane it +// also releases the claim and closes the bulk session. A token that is not the +// one that Attach gave the current holder does nothing. That is how a holder +// that was replaced is stopped from removing the attachment that replaced it. +// +// The token is matched by equality and never by order. A claim made after a +// release can draw a lower number than one that was already issued, so a stale +// token can compare either way against a live one. +// +// To release a claim that this replica no longer holds is normal. It is what a +// worker that moved to another replica looks like from here. Detach is the last +// thing that a dying tunnel does and has nobody to return an error to, so it +// logs. +func (t *Registry) Detach(nodeID string, lane Lane, token int64) { + if lane == LaneBulk { + t.mu.Lock() + held, ok := t.tunnels[nodeID] + if ok && held.bulk != nil && held.bulkToken == token { + held.bulk = nil + held.bulkToken = 0 + } + t.mu.Unlock() + return + } + + t.mu.Lock() + held, ok := t.tunnels[nodeID] + if !ok || held.token != token { + t.mu.Unlock() + return + } + delete(t.tunnels, nodeID) + claim := held.claim + bulk := held.bulk + t.mu.Unlock() + + // The bulk lane lives and ends with the tunnel that it belongs to. Its + // worker dials it again after the inference lane is back. + if bulk != nil { + _ = bulk.Close() + } + + // Not the context of the caller, and not the one that Attach got. Both + // belong to a request or a process that set the tunnel up, and either can be + // cancelled by now, which would leave the row behind on every ordinary + // disconnect. + ctx, cancel := context.WithTimeout(context.Background(), releaseTimeout) + defer cancel() + // The claim, not the token. The row carries the last epoch that was claimed + // for this attachment, and the release must match the row exactly. + if err := t.reg.Release(ctx, nodeID, t.selfID, claim); err != nil { + if errors.Is(err, cluster.ErrNoConnection) { + xlog.Debug("worker tunnel claim was already replaced", "node", nodeID, "epoch", claim) + return + } + xlog.Warn("Releasing a worker tunnel claim failed; peers drop it when the heartbeat of this replica ages out", + "node", nodeID, "epoch", claim, "error", err) + } +} + +// Open returns a stream to the worker over the tunnel that this replica holds. +// +// The bulk lane falls back to the inference lane when the node has no bulk +// session. A worker that predates the bulk lane never has one, and a worker +// that is dialling it again has none for a moment. A transfer that works on the +// session that is shared is better than a transfer that fails. +// +// ErrNotOwner means only that no tunnel of the node is held here. Any other +// failure is returned as itself, wrapped: a session that ended under a held +// entry is a condition of the transport, and ErrNotOwner for it would send a +// dialer to look elsewhere for a worker that this replica holds. +// +// A failed open does not remove the entry. Attach and Detach decide what is +// held here, and an open that removed it would race with the goroutine that +// owns the session and is about to detach it. +func (t *Registry) Open(ctx context.Context, nodeID string, lane Lane) (net.Conn, error) { + t.mu.Lock() + held, ok := t.tunnels[nodeID] + var sess *Session + if ok { + sess = held.inference + if lane == LaneBulk && held.bulk != nil && !held.bulk.IsClosed() { + sess = held.bulk + } + } + t.mu.Unlock() + if sess == nil { + return nil, fmt.Errorf("opening a stream to node %q: %w", nodeID, ErrNotOwner) + } + + stream, err := sess.OpenStream(ctx) + if err != nil { + // A caller whose own budget ran out gets the error of the socket before + // the cancel function of its context has necessarily run. The error + // would then be reported as a tunnel that will not carry a stream, for a + // worker that is well and a client that was impatient. callerRanOut + // decides on the wall clock. The error of the caller is wrapped and not + // returned bare, so that context.DeadlineExceeded can still be matched. + if ctxErr := callerRanOut(ctx); ctxErr != nil { + return nil, fmt.Errorf("opening a stream to node %q over the tunnel held here: the budget of the caller ran out: %w", nodeID, ctxErr) + } + return nil, fmt.Errorf("opening a stream to node %q over the tunnel held here: %w", nodeID, err) + } + return stream, nil +} + +// callerRanOut reports whether the budget of the caller is what ended an +// attempt. +// +// ctx.Err() alone is not that question. A dial passes the deadline of the caller +// down to the socket, so when the budget runs out the timer of the socket fires +// and its error travels up through the multiplexer. The cancellation of the +// context is another timer, and its function has to be run by the scheduler +// before ctx.Err() stops returning nil. Nothing orders the two, so a worker that +// is well could be reported as unreachable to a caller that only ran out of +// time. The wall clock settles it without waiting for a goroutine. +func callerRanOut(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + // A failure at exactly the deadline is the caller's too. + if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) { + return context.DeadlineExceeded + } + return nil +} + +// Holds reports whether this replica holds the inference session of a node that +// has not ended. The connect endpoint asks it before it upgrades a bulk +// session, because after the upgrade only a closed socket can say no. +func (t *Registry) Holds(nodeID string) bool { + t.mu.Lock() + defer t.mu.Unlock() + held, ok := t.tunnels[nodeID] + return ok && !held.inference.IsClosed() +} + +// Held returns the nodes whose tunnels this replica holds, sorted. It answers +// what this process holds, which is another question than who the table says +// owns a node. That question is Owner. +func (t *Registry) Held() []string { + t.mu.Lock() + defer t.mu.Unlock() + out := make([]string, 0, len(t.tunnels)) + for nodeID := range t.tunnels { + out = append(out, nodeID) + } + sort.Strings(out) + return out +} + +// Reclaim writes a new claim for every tunnel that is still held here, and +// returns how many it wrote. +// +// It is for one case. This replica stalled long enough for a peer to sweep it, +// which deleted its instance row and every connection row it owned, and it has +// registered again. Registering again rebuilds the instance row and nothing +// else. Without Reclaim the sockets are held here while the table says that +// nobody holds them, and the other replicas answer "not connected" for workers +// that are connected. +// +// A session that is closed is skipped. A claim is an upsert and takes the row +// from whoever holds it now, and a worker with a closed socket here has already +// connected somewhere else. The entry of a skipped tunnel stays: whoever +// attached it owns its life, and Reclaim is not a way to remove it. +// +// A failure for one node does not stop the others. The next sweep that this +// replica survives tries again. +func (t *Registry) Reclaim(ctx context.Context) (int, error) { + t.mu.Lock() + held := make([]string, 0, len(t.tunnels)) + for nodeID := range t.tunnels { + held = append(held, nodeID) + } + t.mu.Unlock() + + var reclaimed int + var errs []error + for _, nodeID := range held { + if err := t.reclaimOne(ctx, nodeID); err != nil { + if errors.Is(err, errNotReclaimed) { + continue + } + errs = append(errs, err) + continue + } + reclaimed++ + } + if len(errs) > 0 { + return reclaimed, fmt.Errorf("claiming worker tunnels again: %w", errors.Join(errs...)) + } + return reclaimed, nil +} + +// errNotReclaimed means that a node was passed over and did not fail: its +// session is closed, or the attachment went away during the claim. +var errNotReclaimed = errors.New("tunnel: not claimed again") + +// reclaimOne writes a new claim for one node and records it on the attachment +// that is installed for the node. +// +// The attachment that is installed when the gate is taken is the one that gets +// the claim, and not the one that Reclaim listed a moment before. A worker that +// dialled again in between has an entry with its own claim, which the gate makes +// older than this one. The row now holds the new epoch and only this entry can +// release it. +// +// If nothing is installed, the attachment was detached during the claim. Detach +// does not use the gate, so this can happen, and it is the one case where a +// claim is drawn that no attachment will release. The claim is released again. +// That cannot take the row of someone else, because the release matches the +// epoch exactly and no epoch is issued twice. +func (t *Registry) reclaimOne(ctx context.Context, nodeID string) error { + // The gate is taken before the entry is read, so that all that this function + // decides is decided about the attachment that the claim lands on. + if err := t.enterClaim(ctx, nodeID); err != nil { + return fmt.Errorf("claiming node %q again: %w", nodeID, err) + } + + var epoch int64 + var installed bool + if err := func() error { + defer t.leaveClaim(nodeID) + + t.mu.Lock() + tunnel, ok := t.tunnels[nodeID] + t.mu.Unlock() + if !ok { + return errNotReclaimed + } + if tunnel.inference.IsClosed() { + xlog.Debug("not claiming a worker tunnel again: its session is closed", "node", nodeID) + return errNotReclaimed + } + + var err error + epoch, err = t.reg.Claim(ctx, nodeID, t.selfID) + if err != nil { + return err + } + + t.mu.Lock() + current, present := t.tunnels[nodeID] + installed = present + if installed { + // current is the entry that was read above: the gate is held, and + // only Attach and this function install an entry. The entry can + // be gone, because Detach does not use the gate. + current.claim = epoch + } + t.mu.Unlock() + return nil + }(); err != nil { + return err + } + + if installed { + return nil + } + + // Released outside the gate: it is a second round trip, and a worker that + // dials this node again must not wait behind a clean-up. + if err := t.reg.Release(ctx, nodeID, t.selfID, epoch); err != nil && !errors.Is(err, cluster.ErrNoConnection) { + return fmt.Errorf("releasing a new claim of detached node %q: %w", nodeID, err) + } + return errNotReclaimed +} diff --git a/core/services/tunnel/registry_test.go b/core/services/tunnel/registry_test.go new file mode 100644 index 000000000000..0e8a8224d0e5 --- /dev/null +++ b/core/services/tunnel/registry_test.go @@ -0,0 +1,414 @@ +package tunnel_test + +import ( + "context" + "errors" + "io" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("Registry", func() { + var ( + ctx context.Context + db *gorm.DB + clusterR *cluster.Registry + registry *tunnel.Registry + ) + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + Expect(clusterR.Register(ctx, "replica-b", "test", 0, "")).To(Succeed()) + registry = tunnel.NewRegistry(clusterR, "replica-a") + }) + + owner := func(node string) (string, error) { + id, _, err := clusterR.Owner(ctx, node) + return id, err + } + + Describe("the inference lane", func() { + It("claims the node for this replica and holds the session", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + + token, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + + id, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal("replica-a")) + Expect(epoch).To(Equal(token)) + Expect(registry.Held()).To(Equal([]string{"w1"})) + }) + + It("refuses a nil session and an unknown lane without touching the table", func() { + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, nil) + Expect(err).To(HaveOccurred()) + frontend, _ := sessionPair(tunnel.LaneInference) + _, err = registry.Attach(ctx, "w1", tunnel.Lane("express"), frontend) + Expect(err).To(HaveOccurred()) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + Expect(registry.Held()).To(BeEmpty()) + }) + + It("releases the claim and drops the entry on Detach", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + token, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + + registry.Detach("w1", tunnel.LaneInference, token) + + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + Expect(registry.Held()).To(BeEmpty()) + }) + + It("ignores a Detach with a token that is not the live one", func() { + first, _ := sessionPair(tunnel.LaneInference) + second, _ := sessionPair(tunnel.LaneInference) + stale, err := registry.Attach(ctx, "w1", tunnel.LaneInference, first) + Expect(err).ToNot(HaveOccurred()) + live, err := registry.Attach(ctx, "w1", tunnel.LaneInference, second) + Expect(err).ToNot(HaveOccurred()) + Expect(live).ToNot(Equal(stale)) + + registry.Detach("w1", tunnel.LaneInference, stale) + + id, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal("replica-a")) + Expect(epoch).To(Equal(live)) + Expect(registry.Held()).To(Equal([]string{"w1"})) + }) + + It("closes the session that a new dial replaced, and keeps the new one", func() { + first, firstWorker := sessionPair(tunnel.LaneInference) + second, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, first) + Expect(err).ToNot(HaveOccurred()) + _, err = registry.Attach(ctx, "w1", tunnel.LaneInference, second) + Expect(err).ToNot(HaveOccurred()) + + Eventually(first.CloseChan()).Should(BeClosed()) + Eventually(firstWorker.CloseChan()).Should(BeClosed()) + Expect(second.IsClosed()).To(BeFalse()) + }) + + It("does not close a session that is attached twice", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + _, err = registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + Expect(frontend.IsClosed()).To(BeFalse()) + }) + + It("leaves this replica as it was when the claim fails", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + canceled, cancel := context.WithCancel(ctx) + cancel() + + _, err := registry.Attach(canceled, "w1", tunnel.LaneInference, frontend) + Expect(err).To(HaveOccurred()) + Expect(registry.Held()).To(BeEmpty()) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("keeps the entry whose claim the row carries when two dials race", func() { + const dials = 8 + var wg sync.WaitGroup + tokens := make([]int64, dials) + sessions := make([]*tunnel.Session, dials) + for i := range sessions { + sessions[i], _ = sessionPair(tunnel.LaneInference) + } + for i := range dials { + wg.Add(1) + go func() { + defer wg.Done() + defer GinkgoRecover() + tok, err := registry.Attach(ctx, "w1", tunnel.LaneInference, sessions[i]) + Expect(err).ToNot(HaveOccurred()) + tokens[i] = tok + }() + } + wg.Wait() + + _, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + // Exactly one session survives, and it is the one whose claim won + // the row, so its Detach releases the row. + alive := -1 + for i, s := range sessions { + if !s.IsClosed() { + Expect(alive).To(Equal(-1), "two sessions are still open") + alive = i + } + } + Expect(alive).ToNot(Equal(-1)) + Expect(tokens[alive]).To(Equal(epoch)) + + registry.Detach("w1", tunnel.LaneInference, tokens[alive]) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + }) + + Describe("Open", func() { + It("reports ErrNotOwner for a node that is not held here", func() { + _, err := registry.Open(ctx, "ghost", tunnel.LaneInference) + Expect(err).To(MatchError(tunnel.ErrNotOwner)) + }) + + It("opens a stream that reaches the worker", func() { + frontend, worker := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + serveStreams(worker, echoTCP) + + st, err := registry.Open(ctx, "w1", tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 4) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("ping")) + Expect(st.Close()).To(Succeed()) + }) + + It("reports a session that ended under a held entry as itself and not as ErrNotOwner", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + Expect(frontend.Close()).To(Succeed()) + + _, err = registry.Open(ctx, "w1", tunnel.LaneInference) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse()) + Expect(registry.Held()).To(Equal([]string{"w1"}), "a failed open must not remove the entry") + }) + + It("blames the caller and not the tunnel when the budget of the caller is spent", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + Expect(frontend.Close()).To(Succeed()) + + spent, cancel := context.WithDeadline(ctx, time.Now().Add(-time.Second)) + defer cancel() + _, err = registry.Open(spent, "w1", tunnel.LaneInference) + Expect(err).To(MatchError(context.DeadlineExceeded)) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse()) + }) + }) + + Describe("the bulk lane", func() { + attachBoth := func() (inference, bulk, inferenceWorker, bulkWorker *tunnel.Session, infToken, bulkToken int64) { + GinkgoHelper() + inference, inferenceWorker = sessionPair(tunnel.LaneInference) + bulk, bulkWorker = sessionPair(tunnel.LaneBulk) + var err error + infToken, err = registry.Attach(ctx, "w1", tunnel.LaneInference, inference) + Expect(err).ToNot(HaveOccurred()) + bulkToken, err = registry.Attach(ctx, "w1", tunnel.LaneBulk, bulk) + Expect(err).ToNot(HaveOccurred()) + return + } + + It("is held under the claim of the inference lane, without a claim of its own", func() { + _, _, _, _, infToken, _ := attachBoth() + _, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(epoch).To(Equal(infToken)) + }) + + It("is refused with ErrNotOwner when this replica does not hold the inference lane", func() { + bulk, _ := sessionPair(tunnel.LaneBulk) + _, err := registry.Attach(ctx, "w1", tunnel.LaneBulk, bulk) + Expect(err).To(MatchError(tunnel.ErrNotOwner)) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), "a bulk session must not claim the node") + }) + + It("is refused when the inference session has ended", func() { + inference, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, inference) + Expect(err).ToNot(HaveOccurred()) + Expect(inference.Close()).To(Succeed()) + bulk, _ := sessionPair(tunnel.LaneBulk) + _, err = registry.Attach(ctx, "w1", tunnel.LaneBulk, bulk) + Expect(err).To(MatchError(tunnel.ErrNotOwner)) + }) + + It("opens a stream on the lane that is asked for", func() { + _, _, inferenceWorker, bulkWorker, _, _ := attachBoth() + accepted := make(chan string, 2) + serve := func(name string, s *tunnel.Session) { + go func() { + for { + st, err := s.AcceptStream() + if err != nil { + return + } + accepted <- name + _ = st.Close() + } + }() + } + serve("inference", inferenceWorker) + serve("bulk", bulkWorker) + + st, err := registry.Open(ctx, "w1", tunnel.LaneBulk) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("x")) + Expect(err).ToNot(HaveOccurred()) + Eventually(accepted).Should(Receive(Equal("bulk"))) + + st, err = registry.Open(ctx, "w1", tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("x")) + Expect(err).ToNot(HaveOccurred()) + Eventually(accepted).Should(Receive(Equal("inference"))) + }) + + It("falls back to the inference lane when the node has no bulk session", func() { + inference, worker := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, inference) + Expect(err).ToNot(HaveOccurred()) + serveStreams(worker, echoTCP) + + st, err := registry.Open(ctx, "w1", tunnel.LaneBulk) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("ok")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 2) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + }) + + It("falls back to the inference lane when the bulk session has ended", func() { + _, bulk, inferenceWorker, _, _, _ := attachBoth() + serveStreams(inferenceWorker, echoTCP) + Expect(bulk.Close()).To(Succeed()) + + st, err := registry.Open(ctx, "w1", tunnel.LaneBulk) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("ok")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 2) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + }) + + It("closes the session that a new bulk dial replaced", func() { + _, first, _, _, _, _ := attachBoth() + second, _ := sessionPair(tunnel.LaneBulk) + _, err := registry.Attach(ctx, "w1", tunnel.LaneBulk, second) + Expect(err).ToNot(HaveOccurred()) + Eventually(first.CloseChan()).Should(BeClosed()) + Expect(second.IsClosed()).To(BeFalse()) + }) + + It("leaves the inference lane and the claim alone when the bulk lane is detached", func() { + inference, _, _, _, infToken, bulkToken := attachBoth() + registry.Detach("w1", tunnel.LaneBulk, bulkToken) + + Expect(inference.IsClosed()).To(BeFalse()) + _, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(epoch).To(Equal(infToken)) + }) + + It("ignores a bulk Detach with the token of an attachment that was replaced", func() { + _, _, _, _, _, staleBulk := attachBoth() + second, _ := sessionPair(tunnel.LaneBulk) + _, err := registry.Attach(ctx, "w1", tunnel.LaneBulk, second) + Expect(err).ToNot(HaveOccurred()) + + registry.Detach("w1", tunnel.LaneBulk, staleBulk) + Expect(second.IsClosed()).To(BeFalse()) + + st, err := registry.Open(ctx, "w1", tunnel.LaneBulk) + Expect(err).ToNot(HaveOccurred()) + Expect(st.Close()).To(Succeed()) + }) + + It("closes the bulk session with the inference lane that it belongs to", func() { + _, bulk, _, _, infToken, _ := attachBoth() + registry.Detach("w1", tunnel.LaneInference, infToken) + Eventually(bulk.CloseChan()).Should(BeClosed()) + }) + + It("closes the bulk session of an attachment that a new inference dial replaced", func() { + _, bulk, _, _, _, _ := attachBoth() + again, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, again) + Expect(err).ToNot(HaveOccurred()) + Eventually(bulk.CloseChan()).Should(BeClosed()) + }) + }) + + Describe("Reclaim", func() { + It("claims the held tunnels again after this replica was swept", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + // A peer swept this replica: the rows are gone. + Expect(db.Exec("DELETE FROM node_connections").Error).To(Succeed()) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + + n, err := registry.Reclaim(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(Equal(1)) + id, _, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal("replica-a")) + }) + + It("lets the Detach of the attachment release the claim that Reclaim wrote", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + token, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + _, err = registry.Reclaim(ctx) + Expect(err).ToNot(HaveOccurred()) + + registry.Detach("w1", tunnel.LaneInference, token) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("skips a tunnel whose session is closed instead of taking the node back", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + // The worker went to the other replica. + epochB, err := clusterR.Claim(ctx, "w1", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(frontend.Close()).To(Succeed()) + + n, err := registry.Reclaim(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeZero()) + id, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal("replica-b")) + Expect(epoch).To(Equal(epochB)) + }) + }) +}) From 48b50b2c90be2dfe39ee61e786239471a1e658b9 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:38:53 +0000 Subject: [PATCH 021/162] feat(cluster): claim the held connections again after a sweep When a peer sweeps a replica that stalled, it deletes the instance row of that replica and records a departure for every connection it owned. The replica registers again and rebuilds the instance row, but the sockets are still open on it. The other replicas then answer "not connected" for workers that are connected. The membership loop now takes a Reclaimer and calls it after a successful registration that follows a sweep. The interface keeps the package free of the code that holds the sockets. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/membership.go | 53 ++++++++++++++++++ core/services/cluster/membership_test.go | 70 ++++++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index cbec21219f3b..6ea2662843b1 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -66,6 +66,14 @@ func DepartedRetentionFor(grace time.Duration) time.Duration { return DepartedRetention } +// Reclaimer writes the claims of the worker connections that this replica still +// holds. The registry of tunnels is one. The interface keeps this package a +// leaf: it cannot import the package that holds the sockets. +type Reclaimer interface { + // Reclaim returns how many connections it claimed again. + Reclaim(ctx context.Context) (int, error) +} + // Membership publishes this replica in the instances table and removes the // replicas that have stopped answering. // @@ -94,6 +102,7 @@ type Membership struct { retention time.Duration readyEpoch int64 readyReason string + reclaimer Reclaimer } // NewMembership returns the membership loop for one replica. @@ -125,6 +134,23 @@ func (m *Membership) SetReconnectGrace(grace time.Duration) { m.retention = DepartedRetentionFor(grace) } +// SetReclaimer gives the loop the holder of the worker connections of this +// replica. After a sweep deleted the row of this replica, the loop registers it +// again and asks the holder to claim its connections again. Without this, the +// sockets are held here while the table says that nobody holds them. +// +// It is a setter and not an argument of NewMembership because the holder is +// built after the loop. It is safe to call while the loop runs, and on a nil +// receiver. +func (m *Membership) SetReclaimer(r Reclaimer) { + if m == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + m.reclaimer = r +} + // ReportReady records in the instances table that this replica has built the // carrier for the cluster carrier row at epoch, or, with a reason, that it // could not. The loop writes the same values again whenever it has to register @@ -228,7 +254,13 @@ func (m *Membership) tick(ctx context.Context) { // row has to be rebuilt. xlog.Warn("Cluster instance row was reaped, registering again", "id", m.id) if err := m.register(ctx); err != nil { + // The connections are not claimed again in this case. A claim + // written now would name an instance row that does not exist, and + // the next sweep would delete it as an orphan. The sweep below + // still runs, because it removes other replicas. xlog.Error("Registering the cluster instance again failed", "id", m.id, "error", err) + } else { + m.reclaim(ctx) } } else if err != nil { xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) @@ -261,6 +293,27 @@ func (m *Membership) tick(ctx context.Context) { } } +// reclaim claims the connections of this replica again, after the sweep that +// recorded them as departed. The registration that came before it rebuilt the +// instance row only. The lock is not held during the database work. +func (m *Membership) reclaim(ctx context.Context) { + m.mu.Lock() + reclaimer := m.reclaimer + m.mu.Unlock() + if reclaimer == nil { + return + } + n, err := reclaimer.Reclaim(ctx) + if err != nil { + // The loop goes on. The next heartbeat fails in the same way while the + // row is missing, so the claim is tried again. + xlog.Error("Claiming worker connections again after this replica was reaped failed", "id", m.id, "error", err) + } + if n > 0 { + xlog.Info("Claimed worker connections again after this replica was reaped", "id", m.id, "connections", n) + } +} + // Deregister removes one replica and records a departure for every connection // that it owned. // diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go index ee1c90b7e84d..6c4c9de0f153 100644 --- a/core/services/cluster/membership_test.go +++ b/core/services/cluster/membership_test.go @@ -2,6 +2,7 @@ package cluster_test import ( "context" + "sync" "time" "github.com/mudler/LocalAI/core/services/cluster" @@ -425,3 +426,72 @@ var _ = Describe("Reporting readiness through the membership loop", func() { Expect(epoch).To(BeZero()) }) }) + +// countingReclaimer records how often the loop asks it to claim again. +type countingReclaimer struct { + mu sync.Mutex + calls int + n int +} + +func (c *countingReclaimer) Reclaim(context.Context) (int, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.calls++ + return c.n, nil +} + +func (c *countingReclaimer) count() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.calls +} + +var _ = Describe("Claiming connections again after a sweep", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + instanceRows := func(id string) func() int64 { + return func() int64 { + var n int64 + Expect(db.Model(&cluster.Instance{}).Where("id = ?", id).Count(&n).Error).To(Succeed()) + return n + } + } + + It("asks the holder once the row of the replica is back", func() { + membership := cluster.NewMembership(reg, "me", "v1") + reclaimer := &countingReclaimer{n: 2} + membership.SetReclaimer(reclaimer) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + Consistently(reclaimer.count, 2*time.Second).Should(BeZero(), "a normal heartbeat must not claim anything") + + // A peer swept the row while this process stalled. + Expect(reg.Deregister(ctx, "me")).To(Succeed()) + + Eventually(instanceRows("me"), 3*cluster.InstanceHeartbeat, 250*time.Millisecond).Should(Equal(int64(1))) + Eventually(reclaimer.count, time.Second, 50*time.Millisecond).Should(Equal(1)) + }) + + It("works with no holder and with a nil membership", func() { + var nilMembership *cluster.Membership + nilMembership.SetReclaimer(&countingReclaimer{}) + + membership := cluster.NewMembership(reg, "me", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + Expect(reg.Deregister(ctx, "me")).To(Succeed()) + Eventually(instanceRows("me"), 3*cluster.InstanceHeartbeat, 250*time.Millisecond).Should(Equal(int64(1))) + }) +}) From 2088ea31ca357e0e6a2aa6df567f1bc5d82b49a0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:39:06 +0000 Subject: [PATCH 022/162] feat(nodes): mint a tunnel credential for every node at registration A node gets its own secret for the tunnel. Registration mints a new random token with 128 bits or more, returns it once, and stores only its SHA-256 in a new column. The registration token of the deployment never opens a tunnel, so a leak of it does not let its holder be any worker. Every registration mints a new token, and the worker reads it each time it dials. A node waits for approval with a token that the connect route refuses with 403 until the approval. A node type that holds no credential has the column cleared. The answer to a registration now names the active carrier and its epoch when the endpoint has a carrier reader, and hands over the credential for that carrier only: the tunnel token when the tunnel is active, the NATS user JWT otherwise. Without a reader, or when the row cannot be read, the answer is the one that a deployment with only NATS always got. NATS still requires the address of a backend worker. With the tunnel active the address is optional, and a worker that says that its address is not routable has it dropped and cleared. When the tunnel is active and registration is open, the frontend logs one warning. It does not refuse. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/localai/nodes.go | 170 +++++++++++- .../endpoints/localai/nodes_carrier_test.go | 257 ++++++++++++++++++ core/http/routes/nodes.go | 8 +- core/services/nodes/registry.go | 39 +++ 4 files changed, 463 insertions(+), 11 deletions(-) create mode 100644 core/http/endpoints/localai/nodes_carrier_test.go diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index b38cb5a6033f..cf9085c422fd 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -2,6 +2,7 @@ package localai import ( "context" + "crypto/rand" "crypto/sha256" "crypto/subtle" "encoding/hex" @@ -24,6 +25,7 @@ import ( "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" @@ -114,12 +116,64 @@ type RegisterNodeRequest struct { Version string `json:"version,omitempty"` // Commit is the git commit hash the worker binary was built from. Commit string `json:"commit,omitempty"` + // Routable says if the address of the worker can be reached from the + // frontends. A worker started with an address sends true. A worker that has + // none sends false: it holds an outbound tunnel and nothing dials it. A + // worker that predates the field sends nothing, and that means true. + Routable *bool `json:"routable,omitempty"` +} + +// CarrierReader reads the row that names the carrier of the cluster. +// *cluster.CarrierStore is one. +type CarrierReader interface { + Get(ctx context.Context) (cluster.CarrierRow, error) +} + +type registerOptions struct { + carriers CarrierReader +} + +// RegisterOption changes how a registration or an approval answers. +type RegisterOption func(*registerOptions) + +// WithCarrierReader makes the answer name the active carrier and its epoch, and +// hand over the credential for that carrier. Without it the answer is the one of +// a deployment that has only NATS. +func WithCarrierReader(r CarrierReader) RegisterOption { + return func(o *registerOptions) { o.carriers = r } +} + +func newRegisterOptions(opts []RegisterOption) registerOptions { + var o registerOptions + for _, opt := range opts { + opt(&o) + } + return o +} + +// activeCarrier returns the row of the carrier. The second result is false when +// there is no reader or the read failed. The caller then answers as a +// deployment with only NATS does. A failed read is logged and does not fail the +// registration: registration is how a worker joins the cluster, and a worker +// that is refused for a problem it cannot fix would not join at all. +func (o registerOptions) activeCarrier(ctx context.Context) (cluster.CarrierRow, bool) { + if o.carriers == nil { + return cluster.CarrierRow{}, false + } + row, err := o.carriers.Get(ctx) + if err != nil { + xlog.Warn("Reading the cluster carrier for a registration failed; answering as for NATS", "error", err) + return cluster.CarrierRow{}, false + } + return row, true } // RegisterNodeEndpoint registers a new backend node. // expectedToken is the registration token configured on the frontend (may be empty to disable auth). // autoApprove controls whether new nodes go directly to "healthy" or require admin approval. -func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config) echo.HandlerFunc { +func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) + var openRegistrationWarning sync.Once return func(c echo.Context) error { var req RegisterNodeRequest if err := c.Bind(&req); err != nil { @@ -148,11 +202,23 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au fmt.Sprintf("invalid node_type %q; must be %q or %q", nodeType, nodes.NodeTypeBackend, nodes.NodeTypeAgent))) } - // Backend workers require address; agent workers don't serve gRPC + ctx := c.Request().Context() + carrierRow, haveCarrier := opts.activeCarrier(ctx) + tunnelActive := haveCarrier && carrierRow.Active == cluster.CarrierTunnel + // A worker that said that its address cannot be reached holds an + // outbound tunnel and nothing dials it. Its address is dropped and not + // stored: a dialable-looking endpoint that nothing may dial misleads. + tunnelOnly := tunnelActive && req.Routable != nil && !*req.Routable + if tunnelOnly { + req.Address, req.HTTPAddress = "", "" + } + + // Backend workers require address on NATS, which dials it. Agent workers + // don't serve gRPC. With the tunnel active the address is optional. if req.Name == "" { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name is required")) } - if nodeType == nodes.NodeTypeBackend && req.Address == "" { + if nodeType == nodes.NodeTypeBackend && req.Address == "" && !tunnelActive { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address is required for backend workers")) } if len(req.Name) > 255 { @@ -204,11 +270,17 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au Commit: req.Commit, } - ctx := c.Request().Context() if err := registry.Register(ctx, node, autoApprove); err != nil { xlog.Error("Failed to register node", "name", req.Name, "error", err) return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to register node")) } + if tunnelOnly { + // Register keeps a column when the new value is empty, so a node that + // had addresses from an earlier registration still has them. + if err := registry.ClearAddresses(ctx, node.ID); err != nil { + xlog.Warn("Failed to clear the addresses of a tunnel-only node", "node", node.Name, "error", err) + } + } // Merge worker-supplied labels into the node's existing label set, // then apply auto-labels. Using SetNodeLabels here would have @@ -257,7 +329,14 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } - attachNatsJWT(response, node, natsCfg) + attachCarrierCredential(ctx, response, registry, node, natsCfg, carrierRow, haveCarrier) + if tunnelActive && expectedToken == "" && autoApprove { + openRegistrationWarning.Do(func() { + xlog.Warn("Worker registration is open: there is no registration token and new nodes are approved at once, " + + "so anyone who can reach this frontend can register a node, get a tunnel to it, and receive the requests that are sent to workers. " + + "Set a registration token or turn off automatic approval.") + }) + } return c.JSON(http.StatusCreated, response) } @@ -265,7 +344,8 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au // ApproveNodeEndpoint approves a pending node, setting its status to healthy. // For agent workers, it also provisions an API key so they can call the inference API. -func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config) echo.HandlerFunc { +func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) return func(c echo.Context) error { ctx := c.Request().Context() id := c.Param("id") @@ -295,12 +375,88 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr } } - attachNatsJWT(response, node, natsCfg) + // A node that waited for approval already holds its tunnel credential, + // which was minted when it registered and is inert until now. Only the + // NATS credential is minted at approval, and only when NATS is active. + if row, haveCarrier := opts.activeCarrier(ctx); !haveCarrier || row.Active != cluster.CarrierTunnel { + attachNatsJWT(response, node, natsCfg) + } return c.JSON(http.StatusOK, response) } } +// attachCarrierCredential names the active carrier in the answer and adds the +// credential for it: the tunnel token when the tunnel is active, the NATS user +// JWT otherwise. A deployment with no carrier row is a deployment with NATS and +// gets the answer that it always had. +func attachCarrierCredential(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, natsCfg natsauth.Config, row cluster.CarrierRow, haveCarrier bool) { + if haveCarrier { + response["carrier"] = row.Active + response["carrier_epoch"] = row.Epoch + } + if haveCarrier && row.Active == cluster.CarrierTunnel { + attachTunnelToken(ctx, response, registry, node) + return + } + attachNatsJWT(response, node, natsCfg) +} + +// attachTunnelToken mints a new tunnel credential for the node, stores only its +// hash and puts the secret in the response. +// +// Every registration mints a new one, so a second process that registers under +// the same node name makes the next dial of the first process fail. The worker +// reads the token each time it dials. +// +// A node that waits for approval gets one too. That differs from the API key of +// an agent worker, which is not given before approval because it works at once. +// The tunnel credential does not: the connect endpoint reads the status of the +// node at every dial and refuses a pending node with 403. A worker that +// registers once, and gets no credential until it is approved, could not +// connect when the approval came, because an approval does not make it register +// again. +// +// Only a backend or agent node holds a credential. For any other type the +// column is cleared and not left alone. The connect endpoint does not look at +// the type of the node. It refuses a node with an empty hash, so a cleared +// column is what keeps such a node out. Register skips an empty value, so a +// node that registered again as another type would keep a credential that its +// type is not entitled to. +// +// A failure to mint or to store is logged and the answer goes out without the +// token. Registration is how a worker joins the cluster, and a tunnel problem +// must not become a node that cannot join. The worker sees no tunnel_token, +// says so, and tries at its next registration. +func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) { + if node == nil { + return + } + if !tunnelEligible(node.NodeType) { + if err := registry.SetTunnelTokenHash(ctx, node.ID, ""); err != nil { + xlog.Error("Failed to clear the tunnel credential of a node whose type holds none", + "node", node.Name, "type", node.NodeType, "error", err) + } + return + } + // crypto/rand.Text has at least 128 bits of randomness and no error. + secret := rand.Text() + sum := sha256.Sum256([]byte(secret)) + if err := registry.SetTunnelTokenHash(ctx, node.ID, hex.EncodeToString(sum[:])); err != nil { + xlog.Error("Failed to store a tunnel credential for node", "node", node.Name, "error", err) + return + } + response["tunnel_token"] = secret +} + +// tunnelEligible reports whether a node of this type holds a tunnel credential. +// The minting and the clearing both use it, so they cannot disagree: if only one +// were widened, a type would get a credential and never lose it, or lose it at +// every registration. +func tunnelEligible(nodeType string) bool { + return nodeType == nodes.NodeTypeBackend || nodeType == nodes.NodeTypeAgent +} + // attachNatsJWT adds a per-node NATS user JWT to a register/approve response when minting is enabled. func attachNatsJWT(response map[string]any, node *nodes.BackendNode, natsCfg natsauth.Config) { if !natsCfg.CanMintWorkers() || node == nil || node.Status == nodes.StatusPending { diff --git a/core/http/endpoints/localai/nodes_carrier_test.go b/core/http/endpoints/localai/nodes_carrier_test.go new file mode 100644 index 000000000000..d975d5d59351 --- /dev/null +++ b/core/http/endpoints/localai/nodes_carrier_test.go @@ -0,0 +1,257 @@ +package localai + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/pkg/natsauth" + "github.com/nats-io/nkeys" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fixedCarrier answers with a row that the spec chose. +type fixedCarrier struct { + row cluster.CarrierRow + err error +} + +func (f fixedCarrier) Get(context.Context) (cluster.CarrierRow, error) { return f.row, f.err } + +var _ = Describe("Registration and the active carrier", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + natsCfg natsauth.Config + ) + + tunnelRow := cluster.CarrierRow{Active: cluster.CarrierTunnel, Epoch: 7, State: cluster.StateStable} + natsRow := cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 3, State: cluster.StateStable} + + BeforeEach(func() { + ctx = context.Background() + var err error + registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + + akp, err := nkeys.CreateAccount() + Expect(err).ToNot(HaveOccurred()) + seed, err := akp.Seed() + Expect(err).ToNot(HaveOccurred()) + natsCfg = natsauth.Config{AccountSeed: string(seed)} + }) + + register := func(body string, autoApprove bool, opts ...RegisterOption) (int, map[string]any) { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + Expect(RegisterNodeEndpoint(registry, "", autoApprove, nil, "", natsCfg, opts...)(e.NewContext(req, rec))).To(Succeed()) + var resp map[string]any + if rec.Code == http.StatusCreated { + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + } + return rec.Code, resp + } + + storedHash := func(name string) string { + GinkgoHelper() + node, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return node.TunnelTokenHash + } + + Describe("a deployment with no carrier reader", func() { + It("answers as it always did: no carrier, no tunnel token", func() { + code, resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, true) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp).ToNot(HaveKey("carrier")) + Expect(resp).ToNot(HaveKey("carrier_epoch")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(storedHash("w")).To(BeEmpty()) + }) + }) + + Describe("with NATS active", func() { + reader := WithCarrierReader(fixedCarrier{row: natsRow}) + + It("names the carrier and its epoch, and hands over the NATS credential only", func() { + code, resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["carrier"]).To(Equal("nats")) + Expect(resp["carrier_epoch"]).To(BeEquivalentTo(3)) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(storedHash("w")).To(BeEmpty()) + }) + + It("still requires the address of a backend worker", func() { + code, _ := register(`{"name":"w"}`, true, reader) + Expect(code).To(Equal(http.StatusBadRequest)) + }) + + It("keeps the address of a worker that sends routable false, because NATS dials it", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051","routable":false}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051")) + }) + }) + + Describe("with the tunnel active", func() { + reader := WithCarrierReader(fixedCarrier{row: tunnelRow}) + + It("names the carrier and hands over a tunnel token, and no NATS credential", func() { + code, resp := register(`{"name":"w"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["carrier"]).To(Equal("tunnel")) + Expect(resp["carrier_epoch"]).To(BeEquivalentTo(7)) + Expect(resp).ToNot(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("nats_user_seed")) + Expect(resp["tunnel_token"]).To(BeAssignableToTypeOf("")) + }) + + It("stores only the SHA-256 of the token", func() { + _, resp := register(`{"name":"w"}`, true, reader) + token := resp["tunnel_token"].(string) + Expect(len(token)).To(BeNumerically(">=", 26), "128 bits of randomness need at least 26 base32 characters") + + sum := sha256.Sum256([]byte(token)) + Expect(storedHash("w")).To(Equal(hex.EncodeToString(sum[:]))) + Expect(storedHash("w")).ToNot(ContainSubstring(token)) + }) + + It("mints a new token at every registration and the old token stops matching", func() { + _, first := register(`{"name":"w"}`, true, reader) + firstHash := storedHash("w") + _, second := register(`{"name":"w"}`, true, reader) + + Expect(second["tunnel_token"]).ToNot(Equal(first["tunnel_token"])) + Expect(storedHash("w")).ToNot(Equal(firstHash)) + Expect(second["id"]).To(Equal(first["id"]), "the node keeps its identity") + }) + + It("gives a token to a node that waits for approval, and the token is inert until then", func() { + code, resp := register(`{"name":"w"}`, false, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["status"]).To(Equal(nodes.StatusPending)) + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Status).To(Equal(nodes.StatusPending), "the connect endpoint refuses a pending node") + }) + + It("gives a token to an agent worker too", func() { + code, resp := register(`{"name":"agent","node_type":"agent"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + }) + + It("does not need an address", func() { + code, _ := register(`{"name":"w"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + }) + + It("keeps the address of a worker that has a routable one", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051","http_address":"10.0.0.1:50050","routable":true}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051")) + Expect(node.HTTPAddress).To(Equal("10.0.0.1:50050")) + }) + + It("keeps the address of a worker that predates the routable field", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051")) + }) + + It("drops the addresses of a worker that says nothing can reach it, also from an earlier registration", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051","http_address":"10.0.0.1:50050"}`, true, WithCarrierReader(fixedCarrier{row: natsRow})) + Expect(code).To(Equal(http.StatusCreated)) + + code, _ = register(`{"name":"w","address":"host:50051","http_address":"host:50050","routable":false}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(BeEmpty()) + Expect(node.HTTPAddress).To(BeEmpty()) + }) + }) + + It("answers as for NATS when the carrier row cannot be read", func() { + code, resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, true, + WithCarrierReader(fixedCarrier{err: errors.New("database is down")})) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp).ToNot(HaveKey("carrier")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + }) + + Describe("approval", func() { + approve := func(id string, opts ...RegisterOption) map[string]any { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(ApproveNodeEndpoint(registry, nil, "", natsCfg, opts...)(c)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusOK)) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp + } + + It("mints the NATS credential when NATS is active", func() { + _, reg := register(`{"name":"w","address":"10.0.0.1:50051"}`, false) + resp := approve(reg["id"].(string), WithCarrierReader(fixedCarrier{row: natsRow})) + Expect(resp).To(HaveKey("nats_jwt")) + }) + + It("does not mint a NATS credential when the tunnel is active, and leaves the tunnel token alone", func() { + reader := WithCarrierReader(fixedCarrier{row: tunnelRow}) + _, reg := register(`{"name":"w"}`, false, reader) + before := storedHash("w") + + resp := approve(reg["id"].(string), reader) + + Expect(resp).ToNot(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("tunnel_token"), "the worker holds the token from its registration") + Expect(storedHash("w")).To(Equal(before)) + }) + }) + + Describe("a node type that holds no tunnel credential", func() { + It("has the credential cleared and gets no token", func() { + _, reg := register(`{"name":"w"}`, true, WithCarrierReader(fixedCarrier{row: tunnelRow})) + id := reg["id"].(string) + Expect(storedHash("w")).ToNot(BeEmpty()) + + response := map[string]any{} + attachTunnelToken(ctx, response, registry, &nodes.BackendNode{ID: id, Name: "w", NodeType: "gpu-farm"}) + + Expect(response).ToNot(HaveKey("tunnel_token")) + Expect(storedHash("w")).To(BeEmpty()) + }) + }) +}) diff --git a/core/http/routes/nodes.go b/core/http/routes/nodes.go index 2e55cd10b545..53a49e7f4f2f 100644 --- a/core/http/routes/nodes.go +++ b/core/http/routes/nodes.go @@ -36,7 +36,7 @@ func nodeReadyMiddleware(registry *nodes.NodeRegistry) echo.MiddlewareFunc { // token but do not verify per-node identity. A compromised worker can heartbeat/drain/ // deregister other nodes. Future: issue per-node JWT at registration, validate node // identity on subsequent requests (compare :id param with token subject). -func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, registrationToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config) { +func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, registrationToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config, options ...localai.RegisterOption) { if registry == nil { return } @@ -45,7 +45,7 @@ func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, r tokenAuthMw := nodeTokenAuth(registrationToken) node := e.Group("/api/node", readyMw, tokenAuthMw) - node.POST("/register", localai.RegisterNodeEndpoint(registry, registrationToken, autoApprove, authDB, hmacSecret, natsCfg)) + node.POST("/register", localai.RegisterNodeEndpoint(registry, registrationToken, autoApprove, authDB, hmacSecret, natsCfg, options...)) node.POST("/:id/heartbeat", localai.HeartbeatEndpoint(registry)) node.POST("/:id/drain", localai.DrainNodeEndpoint(registry, nil)) node.POST("/:id/resume", localai.ResumeNodeEndpoint(registry)) @@ -61,7 +61,7 @@ func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, r // backend install path (POST /:id/backends/install). That handler enqueues a // ManagementOp on the gallery channel rather than blocking on a NATS reply, so // the browser gets HTTP 202 + jobID immediately instead of waiting up to 3 minutes. -func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender, galleryService *galleryop.GalleryService, opcache *galleryop.OpCache, appConfig *config.ApplicationConfig, adminMw echo.MiddlewareFunc, authDB *gorm.DB, hmacSecret string, registrationToken string, natsCfg natsauth.Config, workerHTTPDial nodes.WorkerNetDialerFor) { +func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender, galleryService *galleryop.GalleryService, opcache *galleryop.OpCache, appConfig *config.ApplicationConfig, adminMw echo.MiddlewareFunc, authDB *gorm.DB, hmacSecret string, registrationToken string, natsCfg natsauth.Config, workerHTTPDial nodes.WorkerNetDialerFor, options ...localai.RegisterOption) { if registry == nil { return } @@ -85,7 +85,7 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade admin.DELETE("/:id", localai.DeregisterNodeEndpoint(registry, unloader)) admin.POST("/:id/drain", localai.DrainNodeEndpoint(registry, unloader)) admin.POST("/:id/resume", localai.ResumeNodeEndpoint(registry)) - admin.POST("/:id/approve", localai.ApproveNodeEndpoint(registry, authDB, hmacSecret, natsCfg)) + admin.POST("/:id/approve", localai.ApproveNodeEndpoint(registry, authDB, hmacSecret, natsCfg, options...)) // Backend management on workers admin.GET("/:id/backends", localai.ListBackendsOnNodeEndpoint(unloader, registry)) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index f1f71b00bbd0..de97f250d671 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -32,6 +32,17 @@ type BackendNode struct { HTTPAddress string `gorm:"size:255" json:"http_address"` // host:port for HTTP file transfer Status string `gorm:"size:32;default:registering" json:"status"` // registering, healthy, unhealthy, draining, pending TokenHash string `gorm:"size:64" json:"-"` // SHA-256 of registration token + // TunnelTokenHash is the SHA-256 of the own credential of this node for the + // tunnel: the token that it presents at GET /api/cluster/connect. It is not + // the registration token. Registration mints a new random secret for each + // node, returns it once and stores only this hash, so a registration token + // that leaked does not open a tunnel for every node whose ID can be read. + // + // Empty means that no tunnel credential was minted for the node. A node that + // registered before the column existed looks like this, and it cannot open a + // tunnel until it registers again. The column cannot be filled in, because + // the secret exists only in the response that minted it. + TunnelTokenHash string `gorm:"size:64" json:"-"` TotalVRAM uint64 `gorm:"column:total_vram" json:"total_vram"` // Total GPU VRAM in bytes AvailableVRAM uint64 `gorm:"column:available_vram" json:"available_vram"` // Available GPU VRAM in bytes // ReservedVRAM is a soft, in-tick reservation deducted by the scheduler when @@ -1340,6 +1351,34 @@ func (r *NodeRegistry) GetWithExtras(ctx context.Context, nodeID string) (*NodeW }, nil } +// SetTunnelTokenHash stores the hash of the tunnel credential of a node. An +// empty hash clears it. It writes the one column on its own because Register +// updates from a struct, and a struct update skips an empty value, so it could +// not clear a credential. +func (r *NodeRegistry) SetTunnelTokenHash(ctx context.Context, nodeID, hash string) error { + res := r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", nodeID).Update("tunnel_token_hash", hash) + if res.Error != nil { + return fmt.Errorf("storing the tunnel credential of node %s: %w", nodeID, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("storing the tunnel credential of node %s: %w", nodeID, gorm.ErrRecordNotFound) + } + return nil +} + +// ClearAddresses removes the addresses of a node. A worker that holds only a +// tunnel advertises none, and Register does not clear a column when the new +// value is empty, so a node that registered with addresses before would keep +// them. +func (r *NodeRegistry) ClearAddresses(ctx context.Context, nodeID string) error { + err := r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", nodeID). + Updates(map[string]any{"address": "", "http_address": ""}).Error + if err != nil { + return fmt.Errorf("clearing the addresses of node %s: %w", nodeID, err) + } + return nil +} + // GetByName returns a single node by name. func (r *NodeRegistry) GetByName(ctx context.Context, name string) (*BackendNode, error) { var node BackendNode From 5ea6a1f2f04b39bee0ec60dcdeb6b8256b10cbe5 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:39:11 +0000 Subject: [PATCH 023/162] feat(http): serve the connect route of the worker tunnel GET /api/cluster/connect takes the dial of a worker. It reads the bearer token before it checks anything else, so an anonymous dial gets 401. A frontend with no registry answers 503 to a caller that has a token. An unknown node gets 401 and not 404, a lookup that fails gets 500, and a node that waits for approval gets 403. The token is compared in constant time with the hash on the row of the node. The registration token of the deployment is never accepted. The upgrade uses the shared upgrader with 64 KiB buffers. The query parameter lane picks the session. A bulk dial to a replica that does not hold the inference session of the node gets 409 before the upgrade, so the worker can dial again at once and land on the owner. The route is registered on every replica. The global auth middleware lets exactly this path through to its own check. The prefix /api/cluster/ is not delegated. The route coverage test walks the route and a new spec fails if the route is removed. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 16 + core/http/app.go | 10 +- core/http/auth/helpers_test.go | 11 + core/http/auth/middleware_test.go | 14 + core/http/auth/public_routes.go | 13 +- .../endpoints/cluster/cluster_suite_test.go | 13 + core/http/endpoints/cluster/connect.go | 214 +++++++++++ core/http/endpoints/cluster/connect_test.go | 340 ++++++++++++++++++ core/http/route_coverage_test.go | 21 ++ core/http/routes/cluster.go | 20 ++ 10 files changed, 669 insertions(+), 3 deletions(-) create mode 100644 core/http/endpoints/cluster/cluster_suite_test.go create mode 100644 core/http/endpoints/cluster/connect.go create mode 100644 core/http/endpoints/cluster/connect_test.go create mode 100644 core/http/routes/cluster.go diff --git a/core/application/distributed.go b/core/application/distributed.go index bef0ca7c4877..e1cc1c959992 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -17,6 +17,7 @@ import ( "github.com/mudler/LocalAI/core/services/agents" "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging" @@ -67,6 +68,15 @@ type DistributedServices struct { // sweeps the rows of replicas that stopped answering. Membership *cluster.Membership + // Carriers reads the row that names the active carrier. The registration + // of a worker answers from it. + Carriers *cluster.CarrierStore + + // Tunnels holds the sessions that workers opened to this replica. It is + // empty while the cluster runs on NATS: the connect route is registered on + // every replica and refuses every worker that has no tunnel credential. + Tunnels *tunnel.Registry + // active names the carrier set the holders forward to. active *atomic.Pointer[carrier.Set] @@ -545,10 +555,16 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade if err != nil { return nil, err } + tunnels := tunnel.NewRegistry(cluster.NewRegistry(authDB), cfg.Distributed.InstanceID) + // If a peer sweeps this replica while it stalls, the claims of the tunnels + // that it still holds are written again when it registers again. + membership.SetReclaimer(tunnels) success = true return &DistributedServices{ Membership: membership, + Carriers: carrierStore, + Tunnels: tunnels, Broadcaster: broadcaster, WorkQueue: workQueue, AgentControl: carrier.NewAgents(active), diff --git a/core/http/app.go b/core/http/app.go index dd51073bda28..995bca4e0910 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -34,6 +34,7 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/quantization" + "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/xlog" ) @@ -562,16 +563,21 @@ func API(application *application.Application) (*echo.Echo, error) { var registry *nodes.NodeRegistry var remoteUnloader nodes.NodeCommandSender var workerHTTPDial nodes.WorkerNetDialerFor + var tunnels *tunnel.Registry + var registerOpts []localai.RegisterOption if d := application.Distributed(); d != nil { registry = d.Registry workerHTTPDial = d.WorkerHTTPDial + tunnels = d.Tunnels + registerOpts = append(registerOpts, localai.WithCarrierReader(d.Carriers)) if d.Router != nil { remoteUnloader = d.Router.Unloader() } } natsCfg := distCfg.NatsAuthConfig() - routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg) - routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDial) + routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg, registerOpts...) + routes.RegisterClusterRoutes(e, registry, tunnels) + routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDial, registerOpts...) // Distributed SSE routes (job progress + agent events via NATS) if d := application.Distributed(); d != nil { diff --git a/core/http/auth/helpers_test.go b/core/http/auth/helpers_test.go index 1b52d4500149..289b3692d6ee 100644 --- a/core/http/auth/helpers_test.go +++ b/core/http/auth/helpers_test.go @@ -122,6 +122,17 @@ func newNodeSelfServiceTestApp(db *gorm.DB, appConfig *config.ApplicationConfig, return e } +// newClusterTestApp uses the real cluster route registration, with no registry, +// so a request that the global auth middleware lets through reaches a handler +// that answers 503 and one that it stops gets 401. +func newClusterTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { + e := echo.New() + e.Use(auth.Middleware(db, appConfig)) + routes.RegisterClusterRoutes(e, nil, nil) + e.GET("/api/cluster/other", ok) + return e +} + // newAdminTestApp creates an Echo app with admin-protected routes. func newAdminTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { e := echo.New() diff --git a/core/http/auth/middleware_test.go b/core/http/auth/middleware_test.go index bcbba0096193..0fa7ab4bd507 100644 --- a/core/http/auth/middleware_test.go +++ b/core/http/auth/middleware_test.go @@ -442,6 +442,20 @@ var _ = Describe("Auth Middleware", func() { } }) + It("delegates the tunnel dial to its own credential check, and only that path", func() { + app = newClusterTestApp(db, appConfig) + + // A bearer token that is no API key and no session. The middleware + // lets the dial through, and the handler (which has no registry here) + // answers 503. If the middleware stopped it, the answer would be 401. + dial := doRequest(app, http.MethodGet, "/api/cluster/connect?id=w1", withBearerToken("own-tunnel-token")) + Expect(dial.Code).To(Equal(http.StatusServiceUnavailable)) + + // Another route under the same prefix is not delegated. + other := doRequest(app, http.MethodGet, "/api/cluster/other", withBearerToken("own-tunnel-token")) + Expect(other.Code).To(Equal(http.StatusUnauthorized)) + }) + It("lets downstream node auth reject missing and invalid registration tokens", func() { const registrationToken = "node-registration-secret" app = newNodeSelfServiceTestApp(db, appConfig, registrationToken) diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index c2509e9a9f6c..fcef6e0d2281 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -80,8 +80,19 @@ func isPublicRoute(method, path string) bool { return false } +// ClusterConnectPath is the route that a worker dials to open its tunnel. The +// handler checks the own tunnel credential of the node, which is not a user, +// a session or an API key, so the global middleware lets the request through. +// +// It is one exact path and not the prefix /api/cluster/, so that a route that is +// added under that prefix later is covered by the global middleware unless its +// author decides otherwise. The path is written out here because this package +// must not import the tunnel package. The spec of the connect endpoint compares +// it with tunnel.ConnectPath. +const ClusterConnectPath = "/api/cluster/connect" + // usesAlternativeAuthentication identifies requests whose credentials are // validated by route-group middleware instead of the global auth middleware. func usesAlternativeAuthentication(path string) bool { - return strings.HasPrefix(path, "/api/node/") + return strings.HasPrefix(path, "/api/node/") || path == ClusterConnectPath } diff --git a/core/http/endpoints/cluster/cluster_suite_test.go b/core/http/endpoints/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..6127ee8c9100 --- /dev/null +++ b/core/http/endpoints/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCluster(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster endpoints test suite") +} diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go new file mode 100644 index 000000000000..ac0aa383c237 --- /dev/null +++ b/core/http/endpoints/cluster/connect.go @@ -0,0 +1,214 @@ +// Package cluster serves the endpoints that the processes of a cluster call +// among themselves. +package cluster + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "errors" + "net/http" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +// ConnectHandler serves the door that a worker dials. It checks the dial against +// the own credential of the node, upgrades it to a websocket, starts the server +// side of a session on it and gives the session to the tunnel registry. +// +// The worker dials out and never listens, so a worker behind NAT, in another +// cluster or on a laptop needs no inbound port. The worker is the yamux client +// and this side is the server. This side opens the streams. It accepts none: the +// frontend asks and the worker answers, so a stream that the worker opened would +// wait in the accept queue for ever. +// +// A worker holds two sessions. The query parameter lane names the one that a +// dial is for: "inference", which is also the default, or "bulk". +// +// The route is registered on every replica of every deployment, so that the +// route coverage test walks it. A registry that is nil is therefore a real +// answer (503) and not a panic. +// +// It is not in auth.RouteFeatureRegistry. That registry gates a route on the +// features of a user, and here there is no user: the caller is a worker process +// with a machine credential. The global auth middleware runs on this path and +// does not reject it, because usesAlternativeAuthentication names the path as +// one that checks its own credentials. Nothing here reads what the middleware +// set. +func ConnectHandler(registry *nodes.NodeRegistry, tunnels *tunnel.Registry) echo.HandlerFunc { + // The default origin check of gorilla refuses a browser of another origin + // and admits a client with no Origin header, which every worker is. + upgrader := tunnel.NewUpgrader() + + return func(c echo.Context) error { + // All checks come before the upgrade, and their order matters. The + // credential is read first, because a dial with no Authorization header + // is the anonymous case and the route coverage test sends exactly that, + // with no query. It must get 401 and not a 400 for a missing node id. + token, ok := bearerToken(c.Request()) + if !ok { + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // Not 401: a frontend that is not distributed cannot authenticate + // anyone, and "unauthorized" would send the operator to look for a + // problem with a token that does not exist. + if registry == nil || tunnels == nil { + return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled") + } + + nodeID := c.QueryParam("id") + if nodeID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing node id") + } + + node, err := registry.Get(c.Request().Context(), nodeID) + switch { + case errors.Is(err, gorm.ErrRecordNotFound): + // 401 and not 404, so that a caller cannot list node IDs by status. + xlog.Debug("worker tunnel dial named an unknown node", "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + case err != nil: + // A query that failed is neither a rejection nor an absence. To tell + // a worker that its credentials are wrong when the database could not + // be read sends it to register again instead of retrying, and a + // worker that registers again has thrown away the credential that + // its tunnel is keyed by. + xlog.Error("Looking up a worker for its tunnel dial failed", "node", nodeID, "error", err) + return echo.NewHTTPError(http.StatusInternalServerError, "node lookup failed") + } + + // An empty hash and a wrong token are two problems with two fixes. An + // empty hash means that the node has not registered since the tunnel + // credential existed, so it has no secret to check and must register. + // A wrong token is a stale credential, usually from before a rotation. + if node.TunnelTokenHash == "" { + // Debug, because every such worker fails in this way on every dial. + xlog.Debug("refusing a worker tunnel: this node has no tunnel credential, so it has not registered since they were introduced", + "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + if !authorizedWorker(token, node.TunnelTokenHash) { + xlog.Debug("worker tunnel dial presented the wrong token", "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // The credential is right, but the node is not allowed yet: 403, because + // the fix is that an admin approves the node and a different credential + // would not help. A tunnel is a standing pipe into the worker that other + // replicas relay to, so a node that waits for approval does not get one. + // Draining and unhealthy nodes keep their tunnels: draining means finish + // what you have, and a node marked unhealthy for missed heartbeats needs + // the pipe to recover. + if node.Status == nodes.StatusPending { + xlog.Warn("Refusing a worker tunnel: this node is awaiting admin approval", "node", nodeID) + return echo.NewHTTPError(http.StatusForbidden, "node is pending approval") + } + + lane, err := tunnel.ParseLane(c.QueryParam("lane")) + if err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } + + // The bulk lane belongs to the replica that holds the inference lane of + // the node. A load balancer can send the dial to another replica. Only a + // status can say that, and the status is possible before the upgrade, so + // the worker is told to dial again at once. + if lane == tunnel.LaneBulk && !tunnels.Holds(nodeID) { + return echo.NewHTTPError(http.StatusConflict, "this replica does not hold the tunnel of the node") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has written its own failure to the client. + xlog.Debug("worker tunnel upgrade failed", "node", nodeID, "error", err) + return nil + } + + sess, err := tunnel.ServerSession(ws, lane) + if err != nil { + xlog.Error("Worker tunnel session setup failed", "node", nodeID, "error", err) + _ = ws.Close() + return nil + } + + // net/http recovers a panic of this goroutine but does not close the + // connection that was hijacked, and the recovery middleware does not + // either. A panic below would leave the worker with a session that this + // replica has no entry for and will never detach. Attach does database + // work, so this is not far-fetched. The panic goes on, because it is a + // bug and the recovery middleware is what reports it. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() + + // The connection is hijacked, so no status can reach the worker any more. + // A failure is a closed socket, which is what the reconnect loop of the + // worker reads. + attachment, err := tunnels.Attach(c.Request().Context(), nodeID, lane, sess) + if err != nil { + xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "lane", lane, "error", err) + _ = sess.Close() + return nil + } + + xlog.Info("Worker tunnel established", "node", nodeID, "lane", lane, "remote", ws.RemoteAddr().String()) + // The session outlives this handler, so something else has to see it end. + // yamux closes its channel from the receive loop as soon as the + // connection fails, and its keepalive notices a worker that vanished + // without a FIN. + go func() { + <-sess.CloseChan() + // The token that Attach returned. Detach compares it by equality: it + // names this attachment, so a worker that has already dialled again + // is not removed by the end of its predecessor. + tunnels.Detach(nodeID, lane, attachment) + xlog.Debug("worker tunnel closed", "node", nodeID, "lane", lane) + }() + return nil + } +} + +// bearerToken returns the token of an Authorization: Bearer header, and whether +// there was one. A missing credential and a wrong credential are different +// answers on purpose: "no credential" decides the 401 before the upgrade, and it +// must be decidable before anything about the node is known. +func bearerToken(r *http.Request) (string, bool) { + // The scheme is not case-sensitive (RFC 7235). The token is. + const prefix = "Bearer " + header := r.Header.Get("Authorization") + if len(header) < len(prefix) || !strings.EqualFold(header[:len(prefix)], prefix) { + return "", false + } + token := header[len(prefix):] + if token == "" { + return "", false + } + return token, true +} + +// authorizedWorker compares a token with the hash on the row of the node, in +// constant time. +// +// The token is the own credential of the node and not the registration token of +// the deployment. A tunnel is a standing pipe into a worker. A credential that +// opens every worker would let one leak take over the traffic of any worker +// whose ID can be read, by claiming its tunnel. The secret is minted for each +// node at registration, returned to that worker once and stored as this hash. +// It is not TokenHash, which hashes the token that the worker registered with. +// On most deployments that is the shared registration token. +func authorizedWorker(token, storedHash string) bool { + if storedHash == "" { + return false + } + sum := sha256.Sum256([]byte(token)) + return subtle.ConstantTimeCompare([]byte(hex.EncodeToString(sum[:])), []byte(storedHash)) == 1 +} diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go new file mode 100644 index 000000000000..cab45f4c75c6 --- /dev/null +++ b/core/http/endpoints/cluster/connect_test.go @@ -0,0 +1,340 @@ +package cluster_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "net/http" + "net/http/httptest" + "time" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/auth" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +func hashOf(secret string) string { + sum := sha256.Sum256([]byte(secret)) + return hex.EncodeToString(sum[:]) +} + +var _ = Describe("Connect endpoint", func() { + const ( + registrationToken = "shared-registration-token" + nodeToken = "own-tunnel-token-of-w1" + ) + + var ( + ctx context.Context + db *gorm.DB + clusterR *cluster.Registry + nodeReg *nodes.NodeRegistry + tunnels *tunnel.Registry + srv *httptest.Server + nodeID string + ) + + // register adds a node the way the registration endpoint does: the + // registration token is stored as TokenHash, and the own credential of the + // node as TunnelTokenHash. + register := func(name string, autoApprove bool) *nodes.BackendNode { + GinkgoHelper() + node := &nodes.BackendNode{Name: name, NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken)} + Expect(nodeReg.Register(ctx, node, autoApprove)).To(Succeed()) + return node + } + + serve := func(registry *nodes.NodeRegistry, tunnels *tunnel.Registry) *httptest.Server { + GinkgoHelper() + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(registry, tunnels)) + s := httptest.NewServer(e) + DeferCleanup(s.Close) + return s + } + + dial := func(query, token string, header ...http.Header) (*websocket.Conn, *http.Response, error) { + h := http.Header{} + if token != "" { + h.Set("Authorization", "Bearer "+token) + } + for _, extra := range header { + for k, v := range extra { + h[k] = v + } + } + url := "ws" + srv.URL[len("http"):] + tunnel.ConnectPath + query + return tunnel.NewDialer(5*time.Second).Dial(url, h) + } + + // plainGet is a request that is no websocket upgrade. + plainGet := func(query, token string) *http.Response { + GinkgoHelper() + req, err := http.NewRequest(http.MethodGet, srv.URL+tunnel.ConnectPath+query, nil) + Expect(err).ToNot(HaveOccurred()) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + return resp + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + var err error + nodeReg, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + tunnels = tunnel.NewRegistry(clusterR, "replica-a") + + node := register("w1", true) + nodeID = node.ID + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hashOf(nodeToken))).To(Succeed()) + + srv = serve(nodeReg, tunnels) + }) + + Describe("the answers before the upgrade", func() { + It("answers 401 to a dial with no credential, before anything else is checked", func() { + // No query string at all: an anonymous dial must not get a 400. + Expect(plainGet("", "").StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(plainGet("?id="+nodeID, "").StatusCode).To(Equal(http.StatusUnauthorized)) + // Not even the 503 for a missing registry may come first. + bare := serve(nil, nil) + resp, err := http.Get(bare.URL + tunnel.ConnectPath) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 503 to a credential when the frontend has no registry", func() { + bare := serve(nil, nil) + req, _ := http.NewRequest(http.MethodGet, bare.URL+tunnel.ConnectPath+"?id=x", nil) + req.Header.Set("Authorization", "Bearer anything") + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + }) + + It("answers 400 to a credential with no node id", func() { + Expect(plainGet("", nodeToken).StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("answers 401, not 404, for a node it does not know", func() { + Expect(plainGet("?id=nobody", nodeToken).StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 401 to a wrong token", func() { + Expect(plainGet("?id="+nodeID, "not-the-token").StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("never accepts the shared registration token", func() { + Expect(plainGet("?id="+nodeID, registrationToken).StatusCode).To(Equal(http.StatusUnauthorized)) + _, resp, err := dial("?id="+nodeID, registrationToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tunnels.Held()).To(BeEmpty()) + }) + + It("answers 401 to a node that has no tunnel credential, whatever it presents", func() { + other := register("w-old", true) + Expect(plainGet("?id="+other.ID, "").StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(plainGet("?id="+other.ID, hashOf("")).StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(plainGet("?id="+other.ID, registrationToken).StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 403 to a node that waits for approval, and takes no tunnel from it", func() { + pending := register("w-pending", false) + Expect(pending.Status).To(Equal(nodes.StatusPending)) + Expect(nodeReg.SetTunnelTokenHash(ctx, pending.ID, hashOf("pending-token"))).To(Succeed()) + + _, resp, err := dial("?id="+pending.ID, "pending-token") + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) + Expect(tunnels.Held()).To(BeEmpty()) + + Expect(nodeReg.ApproveNode(ctx, pending.ID)).To(Succeed()) + ws, _, err := dial("?id="+pending.ID, "pending-token") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(ws.Close) + }) + + It("answers 500 and not 401 when the node cannot be read", func() { + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + Expect(sqlDB.Close()).To(Succeed()) + Expect(plainGet("?id="+nodeID, nodeToken).StatusCode).To(Equal(http.StatusInternalServerError)) + }) + + It("answers 400 to a lane it does not know", func() { + Expect(plainGet("?id="+nodeID+"&lane=express", nodeToken).StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("leaves the table alone when the request is no upgrade", func() { + resp := plainGet("?id="+nodeID, nodeToken) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + _, _, err := clusterR.Owner(ctx, nodeID) + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("refuses a browser of another origin", func() { + _, resp, err := dial("?id="+nodeID, nodeToken, http.Header{"Origin": []string{"https://evil.example"}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) + Expect(tunnels.Held()).To(BeEmpty()) + }) + }) + + Describe("a dial with the own credential of the node", func() { + It("claims the node for this replica and carries streams to the worker", func() { + ws, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + worker, err := tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(worker.Close) + + Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + owner, _, err := clusterR.Owner(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("replica-a")) + + go func() { + for { + st, err := worker.AcceptStream() + if err != nil { + return + } + go func() { defer func() { _ = st.Close() }(); _, _ = io.Copy(st, st) }() + } + }() + st, err := tunnels.Open(ctx, nodeID, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 4) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("ping")) + }) + + It("releases the claim when the worker goes away", func() { + ws, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + worker, err := tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + + Expect(worker.Close()).To(Succeed()) + + Eventually(tunnels.Held, 5*time.Second).Should(BeEmpty()) + // The registry drops the entry first and releases the claim after it. + Eventually(func() error { + _, _, err := clusterR.Owner(ctx, nodeID) + return err + }, 5*time.Second).Should(MatchError(cluster.ErrNoConnection)) + }) + + It("replaces the tunnel when the worker dials again", func() { + ws1, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + first, err := tunnel.ClientSession(ws1, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + + ws2, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + second, err := tunnel.ClientSession(ws2, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Close) + + Eventually(first.CloseChan(), 5*time.Second).Should(BeClosed()) + Expect(second.IsClosed()).To(BeFalse()) + // The end of the first session must not remove the second. + Consistently(tunnels.Held, time.Second).Should(Equal([]string{nodeID})) + _, _, err = clusterR.Owner(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + }) + + It("stops accepting the old credential after a new one is minted", func() { + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hashOf("rotated"))).To(Succeed()) + _, resp, err := dial("?id="+nodeID, nodeToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + ws, _, err := dial("?id="+nodeID, "rotated") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(ws.Close) + }) + }) + + Describe("the bulk lane", func() { + It("is refused with 409 before the inference lane is held, and accepted after", func() { + _, resp, err := dial("?id="+nodeID+"&lane=bulk", nodeToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusConflict)) + _, _, ownerErr := clusterR.Owner(ctx, nodeID) + Expect(ownerErr).To(MatchError(cluster.ErrNoConnection), "a bulk dial must not claim the node") + + wsI, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + inference, err := tunnel.ClientSession(wsI, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(inference.Close) + Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + + wsB, _, err := dial("?id="+nodeID+"&lane=bulk", nodeToken) + Expect(err).ToNot(HaveOccurred()) + bulk, err := tunnel.ClientSession(wsB, tunnel.LaneBulk) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(bulk.Close) + + // A transfer stream reaches the worker on the bulk session. + reached := make(chan string, 1) + go func() { + st, err := bulk.AcceptStream() + if err != nil { + return + } + buf := make([]byte, 4) + _, _ = io.ReadFull(st, buf) + reached <- string(buf) + }() + Eventually(func() error { + st, err := tunnels.Open(ctx, nodeID, tunnel.LaneBulk) + if err != nil { + return err + } + _, err = st.Write([]byte("blob")) + return err + }, 5*time.Second).Should(Succeed()) + Eventually(reached).Should(Receive(Equal("blob"))) + }) + + It("authenticates the bulk lane like the inference lane", func() { + _, resp, err := dial("?id="+nodeID+"&lane=bulk", registrationToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + }) +}) + +var _ = Describe("The path of the connect route", func() { + It("is the one that the global auth middleware lets through", func() { + Expect(auth.ClusterConnectPath).To(Equal(tunnel.ConnectPath)) + }) +}) diff --git a/core/http/route_coverage_test.go b/core/http/route_coverage_test.go index 65dfe37f497a..7505cb821f2f 100644 --- a/core/http/route_coverage_test.go +++ b/core/http/route_coverage_test.go @@ -15,6 +15,7 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/config" . "github.com/mudler/LocalAI/core/http" + "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/LocalAI/pkg/system" . "github.com/onsi/ginkgo/v2" @@ -201,4 +202,24 @@ var _ = Describe("Route auth coverage", func() { "global auth middleware or RequireAdmin / RequireFeature.") } }) + + // The route that a worker dials to open its tunnel is registered on every + // replica, also when the deployment is not distributed. The generic walk + // above covers it, but only while it exists. This keeps its removal from + // passing unseen, and checks that a dial with no credential gets 401 and not + // the 503 of a frontend with no registry. + It("registers the worker tunnel route and refuses an anonymous dial", func() { + found := false + for _, r := range app.Routes() { + if r.Method == http.MethodGet && r.Path == tunnel.ConnectPath { + found = true + } + } + Expect(found).To(BeTrue(), "GET %s is not registered", tunnel.ConnectPath) + + req := httptest.NewRequest(http.MethodGet, tunnel.ConnectPath, nil) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusUnauthorized)) + }) }) diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go new file mode 100644 index 000000000000..5e1df77513b2 --- /dev/null +++ b/core/http/routes/cluster.go @@ -0,0 +1,20 @@ +package routes + +import ( + "github.com/labstack/echo/v4" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/tunnel" +) + +// RegisterClusterRoutes registers the routes that the processes of a cluster +// call among themselves. They check their own credentials, and the global auth +// middleware lets them through (see auth.ClusterConnectPath). +// +// The connect route is registered on every replica, in every deployment, and +// refuses every worker that has no tunnel credential. A frontend that is not +// distributed answers 503 to a caller that has a credential and 401 to one that +// has none. +func RegisterClusterRoutes(e *echo.Echo, registry *nodes.NodeRegistry, tunnels *tunnel.Registry) { + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(registry, tunnels)) +} From b27be577e318aa71dd81cdee5514ee217e745273 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:50:29 +0000 Subject: [PATCH 024/162] test(tunnel): move the slow link of the latency spec to a package The proxy that works as a link with a rate and a delay is now a package, so that the specs of the worker can use it on the real endpoint. It also passes the end of a connection on behind the data that is still queued. Before, a connection that ended was closed at once and the data in the queue was lost. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/latency_test.go | 165 ++----------------- core/services/tunnel/slowlink/slowlink.go | 185 ++++++++++++++++++++++ 2 files changed, 194 insertions(+), 156 deletions(-) create mode 100644 core/services/tunnel/slowlink/slowlink.go diff --git a/core/services/tunnel/latency_test.go b/core/services/tunnel/latency_test.go index c1c6002f4808..6dd4be37dea7 100644 --- a/core/services/tunnel/latency_test.go +++ b/core/services/tunnel/latency_test.go @@ -7,171 +7,24 @@ import ( "net/http" "net/http/httptest" "slices" - "sync" "time" "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/tunnel/slowlink" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) -// slowLink is a TCP proxy that works as one bottleneck link for every -// connection that goes through it. It has a rate and a delay in each -// direction, and one queue of 256 KiB for each direction that all connections -// share. A real link has that shape, and the effect under test depends on it: a -// probe that has its own connection still waits in the queue behind the bytes -// of a transfer on another connection. What a probe must not wait for are the -// bytes that sit in the buffers of its own connection ahead of it. -type slowLink struct { - lis net.Listener - target string - rate float64 - oneWay time.Duration - up *linkQueue - down *linkQueue - quit chan struct{} - wg sync.WaitGroup - - mu sync.Mutex - conns []net.Conn -} - -type linkItem struct { - data []byte - dst net.Conn - at time.Time -} - -// linkQueue is one direction of the link. -type linkQueue struct { - in chan linkItem - delivery chan linkItem -} - -func newSlowLink(target string, bytesPerSecond float64, oneWay time.Duration) *slowLink { +// newSlowLink starts a link in front of target and closes it after the spec. +func newSlowLink(target string, bytesPerSecond float64, oneWay time.Duration) *slowlink.Link { GinkgoHelper() - lis, err := net.Listen("tcp", "127.0.0.1:0") + l, err := slowlink.New(target, bytesPerSecond, oneWay) Expect(err).ToNot(HaveOccurred()) - l := &slowLink{ - lis: lis, target: target, rate: bytesPerSecond, oneWay: oneWay, - up: &linkQueue{in: make(chan linkItem, 4), delivery: make(chan linkItem, 4096)}, - down: &linkQueue{in: make(chan linkItem, 4), delivery: make(chan linkItem, 4096)}, - quit: make(chan struct{}), - } - for _, q := range []*linkQueue{l.up, l.down} { - l.wg.Add(2) - go l.transmit(q) - go l.deliver(q) - } - go l.accept() - DeferCleanup(l.close) + DeferCleanup(l.Close) return l } -func (l *slowLink) addr() string { return l.lis.Addr().String() } - -func (l *slowLink) close() { - select { - case <-l.quit: - return - default: - } - close(l.quit) - _ = l.lis.Close() - l.mu.Lock() - for _, c := range l.conns { - _ = c.Close() - } - l.mu.Unlock() - l.wg.Wait() -} - -func (l *slowLink) accept() { - for { - c, err := l.lis.Accept() - if err != nil { - return - } - s, err := net.Dial("tcp", l.target) - if err != nil { - _ = c.Close() - continue - } - l.mu.Lock() - l.conns = append(l.conns, c, s) - l.mu.Unlock() - go l.read(c, s, l.up) - go l.read(s, c, l.down) - } -} - -// read moves what arrives on src into the queue of a direction. A full queue -// stops the read, so the socket fills and the sender slows down, as it would -// on a link with a queue of limited size. -func (l *slowLink) read(src, dst net.Conn, q *linkQueue) { - defer func() { _ = dst.Close() }() - for { - buf := make([]byte, 64<<10) - n, err := src.Read(buf) - if n > 0 { - select { - case q.in <- linkItem{data: buf[:n], dst: dst}: - case <-l.quit: - return - } - } - if err != nil { - return - } - } -} - -// transmit takes an item out of the queue at the rate of the link. -func (l *slowLink) transmit(q *linkQueue) { - defer l.wg.Done() - next := time.Now() - for { - var item linkItem - select { - case item = <-q.in: - case <-l.quit: - return - } - // A late wake-up must not slow the link: the schedule catches up - // instead. Only an idle link starts a new schedule. - if now := time.Now(); now.Sub(next) > 5*time.Millisecond { - next = now - } - next = next.Add(time.Duration(float64(len(item.data)) / l.rate * float64(time.Second))) - if d := time.Until(next); d > 500*time.Microsecond { - time.Sleep(d) - } - item.at = time.Now().Add(l.oneWay) - select { - case q.delivery <- item: - case <-l.quit: - return - } - } -} - -// deliver writes an item to its connection after the delay of the link. One -// goroutine for each direction keeps the order of each connection. -func (l *slowLink) deliver(q *linkQueue) { - defer l.wg.Done() - for { - var item linkItem - select { - case item = <-q.delivery: - case <-l.quit: - return - } - time.Sleep(time.Until(item.at)) - _, _ = item.dst.Write(item.data) - } -} - // probeStats are the delays of the probes in milliseconds. type probeStats struct { count int @@ -411,7 +264,7 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma link := newSlowLink(front, linkRate, linkOneWay) dial := func(kind byte) net.Conn { - c, err := net.Dial("tcp", link.addr()) + c, err := net.Dial("tcp", link.Addr()) Expect(err).ToNot(HaveOccurred()) DeferCleanup(c.Close) _, err = c.Write([]byte{kind}) @@ -428,7 +281,7 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma It("measures the probe on the same session as the transfer", func() { f := startLanedFrontend() link := newSlowLink(f.addr, linkRate, linkOneWay) - worker := dialThrough(link.addr(), tunnel.LaneInference) + worker := dialThrough(link.Addr(), tunnel.LaneInference) var frontend *tunnel.Session Eventually(f.inference).Should(Receive(&frontend)) DeferCleanup(func() { _ = frontend.Close() }) @@ -452,8 +305,8 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma // Both sessions go through the same link, as the two websockets of one // worker do. link := newSlowLink(f.addr, linkRate, linkOneWay) - inferenceWorker := dialThrough(link.addr(), tunnel.LaneInference) - bulkWorker := dialThrough(link.addr(), tunnel.LaneBulk) + inferenceWorker := dialThrough(link.Addr(), tunnel.LaneInference) + bulkWorker := dialThrough(link.Addr(), tunnel.LaneBulk) var inferenceFrontend, bulkFrontend *tunnel.Session Eventually(f.inference).Should(Receive(&inferenceFrontend)) Eventually(f.bulk).Should(Receive(&bulkFrontend)) diff --git a/core/services/tunnel/slowlink/slowlink.go b/core/services/tunnel/slowlink/slowlink.go new file mode 100644 index 000000000000..06172b02ad43 --- /dev/null +++ b/core/services/tunnel/slowlink/slowlink.go @@ -0,0 +1,185 @@ +// Package slowlink is a TCP proxy for tests that works as one link with a rate +// and a delay. It lets a spec measure how a small call fares during a large +// transfer on a link that behaves like a real one, which loopback does not. +package slowlink + +import ( + "net" + "sync" + "time" +) + +// Link is a proxy in front of one target. It has a rate and a delay in each +// direction, and one queue of 256 KiB for each direction that every connection +// shares. A real link has that shape, and the effect under test depends on it: a +// probe that has its own connection still waits in the queue behind the bytes of +// a transfer on another connection. What a probe must not wait for are the bytes +// that sit in the buffers of its own connection ahead of it. +type Link struct { + lis net.Listener + target string + rate float64 + oneWay time.Duration + up *queue + down *queue + quit chan struct{} + wg sync.WaitGroup + + mu sync.Mutex + conns []net.Conn +} + +type item struct { + data []byte + dst net.Conn + at time.Time + // end marks the end of the data of a connection: dst is half-closed after + // everything before it was delivered. + end bool +} + +// queue is one direction of the link. +type queue struct { + in chan item + delivery chan item +} + +// New starts a link on a loopback port that forwards to target. The rate is in +// bytes per second and oneWay is the delay in each direction. +func New(target string, bytesPerSecond float64, oneWay time.Duration) (*Link, error) { + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, err + } + l := &Link{ + lis: lis, target: target, rate: bytesPerSecond, oneWay: oneWay, + up: &queue{in: make(chan item, 4), delivery: make(chan item, 4096)}, + down: &queue{in: make(chan item, 4), delivery: make(chan item, 4096)}, + quit: make(chan struct{}), + } + for _, q := range []*queue{l.up, l.down} { + l.wg.Add(2) + go l.transmit(q) + go l.deliver(q) + } + go l.accept() + return l, nil +} + +// Addr returns the address to dial. +func (l *Link) Addr() string { return l.lis.Addr().String() } + +// Close stops the link and closes every connection that goes through it. +func (l *Link) Close() { + select { + case <-l.quit: + return + default: + } + close(l.quit) + _ = l.lis.Close() + l.mu.Lock() + for _, c := range l.conns { + _ = c.Close() + } + l.mu.Unlock() + l.wg.Wait() +} + +func (l *Link) accept() { + for { + c, err := l.lis.Accept() + if err != nil { + return + } + s, err := net.Dial("tcp", l.target) + if err != nil { + _ = c.Close() + continue + } + l.mu.Lock() + l.conns = append(l.conns, c, s) + l.mu.Unlock() + go l.read(c, s, l.up) + go l.read(s, c, l.down) + } +} + +// read moves what arrives on src into the queue of a direction. A full queue +// stops the read, so the socket fills and the sender slows down, as it would on +// a link with a queue of limited size. +func (l *Link) read(src, dst net.Conn, q *queue) { + for { + buf := make([]byte, 64<<10) + n, err := src.Read(buf) + if n > 0 { + select { + case q.in <- item{data: buf[:n], dst: dst}: + case <-l.quit: + return + } + } + if err != nil { + // The end goes through the queue, behind the data. If dst were + // closed here, the data in the queue would be lost. + select { + case q.in <- item{dst: dst, end: true}: + case <-l.quit: + } + return + } + } +} + +// transmit takes an item out of the queue at the rate of the link. +func (l *Link) transmit(q *queue) { + defer l.wg.Done() + next := time.Now() + for { + var it item + select { + case it = <-q.in: + case <-l.quit: + return + } + // A late wake-up must not slow the link: the schedule catches up + // instead. Only an idle link starts a new schedule. + if now := time.Now(); now.Sub(next) > 5*time.Millisecond { + next = now + } + next = next.Add(time.Duration(float64(len(it.data)) / l.rate * float64(time.Second))) + if d := time.Until(next); d > 500*time.Microsecond { + time.Sleep(d) + } + it.at = time.Now().Add(l.oneWay) + select { + case q.delivery <- it: + case <-l.quit: + return + } + } +} + +// deliver writes an item to its connection after the delay of the link. One +// goroutine for each direction keeps the order of each connection. +func (l *Link) deliver(q *queue) { + defer l.wg.Done() + for { + var it item + select { + case it = <-q.delivery: + case <-l.quit: + return + } + time.Sleep(time.Until(it.at)) + if it.end { + if tcp, ok := it.dst.(*net.TCPConn); ok { + _ = tcp.CloseWrite() + } else { + _ = it.dst.Close() + } + continue + } + _, _ = it.dst.Write(it.data) + } +} From 8f9d08692220bbfc50f5109a062f0cb77b2fc558 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:50:29 +0000 Subject: [PATCH 025/162] feat(nodes): add the readiness probe of a worker on the tunnel A worker that is attached to the tunnel is ready when it holds a tunnel session. The probe is the counterpart of the one for NATS, and its error is a different one, so that /readyz says which link is down. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/worker_readiness.go | 32 ++++++++++++++++++++ core/services/nodes/worker_readiness_test.go | 24 +++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/core/services/nodes/worker_readiness.go b/core/services/nodes/worker_readiness.go index a0d5b3a19490..46a047994628 100644 --- a/core/services/nodes/worker_readiness.go +++ b/core/services/nodes/worker_readiness.go @@ -75,6 +75,38 @@ func NATSReadiness(conn natsConn) func() error { } } +// tunnelConn is the slice of the tunnel client of a worker that the readiness +// probe needs. It is a local interface for the same reasons as natsConn. +type tunnelConn interface { + Connected() bool +} + +// ErrTunnelDisconnected is reported by TunnelReadiness when the worker holds no +// tunnel session. +var ErrTunnelDisconnected = errors.New("worker tunnel is down: the frontend cannot reach this worker") + +// TunnelReadiness builds the readiness probe of a worker that is attached to the +// tunnel carrier. It is the counterpart of NATSReadiness. +// +// Everything that a worker does arrives over its tunnel: the installs and stops +// of backends, the model lifecycle, file staging, and every inference stream. A +// worker with no tunnel session is up and cannot be reached, and a worker that +// binds only loopback has no other way in. The registration of the worker is +// implied by the probe being reachable at all, because the file-transfer server +// starts after the registration. +// +// The heartbeat of the node row is an HTTP call to the frontend and does not +// pass through the tunnel. A worker can keep sending it while its tunnel is +// down, so the local probe sees what the registry cannot. +func TunnelReadiness(conn tunnelConn) func() error { + return func() error { + if conn == nil || !conn.Connected() { + return ErrTunnelDisconnected + } + return nil + } +} + // ErrBackendUnreachable is reported when a worker holds a backend process it // can no longer reach. var ErrBackendUnreachable = errors.New("backend process is unreachable") diff --git a/core/services/nodes/worker_readiness_test.go b/core/services/nodes/worker_readiness_test.go index 6771291ced47..b21c0d67a30e 100644 --- a/core/services/nodes/worker_readiness_test.go +++ b/core/services/nodes/worker_readiness_test.go @@ -16,6 +16,11 @@ type fakeConn struct{ connected bool } func (f *fakeConn) IsConnected() bool { return f.connected } +// fakeTunnel stands in for the tunnel client of a worker. +type fakeTunnel struct{ connected bool } + +func (f *fakeTunnel) Connected() bool { return f.connected } + var _ = Describe("WorkerReadiness", func() { Describe("the gate itself", func() { It("reports ready when no probe has been installed yet", func() { @@ -52,6 +57,25 @@ var _ = Describe("WorkerReadiness", func() { }) }) + Describe("TunnelReadiness", func() { + It("reports ready while the tunnel holds a session", func() { + Expect(TunnelReadiness(&fakeTunnel{connected: true})()).To(Succeed()) + }) + + It("reports not ready once the session is gone", func() { + err := TunnelReadiness(&fakeTunnel{connected: false})() + Expect(err).To(MatchError(ErrTunnelDisconnected)) + }) + + It("reports not ready when there is no tunnel at all", func() { + Expect(TunnelReadiness(nil)()).To(MatchError(ErrTunnelDisconnected)) + }) + + It("does not mix up the reasons of the two carriers", func() { + Expect(errors.Is(ErrTunnelDisconnected, ErrNATSDisconnected)).To(BeFalse()) + }) + }) + Describe("the file transfer server probes", func() { var ( srv *http.Server From 67cf45c324300037c7307e4c5fff55d531b1683f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:50:29 +0000 Subject: [PATCH 026/162] feat(worker): keep the credentials of the carrier that registration named The credential manager is no longer only for NATS. The answer to a registration names the carrier and its epoch. When the carrier is the tunnel, the manager waits for the tunnel token and not for the NATS JWT, and it does not wait for the approval of the node: the frontend mints the token at registration and refuses the dial of a node that is still pending, so the tunnel client waits for the approval itself. The token of the latest registration is the one that the client presents, because every registration mints a new one. For NATS, or for a frontend that names no carrier, the manager works as before. RegisterWithRetry now uses RegisterFullWithRetry, which returns the whole answer. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/agent_worker.go | 2 +- core/cli/workerregistry/client.go | 32 ++++++-- core/cli/workerregistry/credentials.go | 83 +++++++++++++++------ core/cli/workerregistry/credentials_test.go | 78 +++++++++++++++++-- 4 files changed, 160 insertions(+), 35 deletions(-) diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 9c5422a51f7b..9eb493314aca 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -126,7 +126,7 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // Acquire credentials via (re)registration. When the bus requires auth and no // static fallback is configured, wait through admin approval until the // frontend mints credentials rather than starting unauthenticated. - credMgr := workerregistry.NewNATSCredentialManager( + credMgr := workerregistry.NewCredentialManager( func(ctx context.Context) (*workerregistry.RegisterResponse, error) { return regClient.RegisterFull(ctx, registrationBody) }, diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index cf46455c95c0..6bec5ca95a01 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -63,6 +63,14 @@ type RegisterResponse struct { APIToken string `json:"api_token,omitempty"` NatsJWT string `json:"nats_jwt,omitempty"` NatsUserSeed string `json:"nats_user_seed,omitempty"` + // Carrier is the carrier that is active in the cluster: "nats" or "tunnel". + // A frontend that predates carriers does not send it, and then the cluster + // runs on NATS. CarrierEpoch is the epoch of the carrier row at that time. + Carrier string `json:"carrier,omitempty"` + CarrierEpoch int64 `json:"carrier_epoch,omitempty"` + // TunnelToken is the own credential of the node for the tunnel. It is sent + // once and only when the tunnel is the active carrier. + TunnelToken string `json:"tunnel_token,omitempty"` } // RegisterFull sends a single registration request and returns the full @@ -109,26 +117,40 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) // RegisterWithRetry retries registration with exponential backoff. func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) { + res, err := c.RegisterFullWithRetry(ctx, body, maxRetries) + if err != nil { + return "", "", "", "", err + } + return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil +} + +// RegisterFullWithRetry is RegisterWithRetry that returns the whole response. +func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map[string]any, maxRetries int) (*RegisterResponse, error) { + if maxRetries < 1 { + return nil, fmt.Errorf("registering: %d attempts allowed, want at least one", maxRetries) + } backoff := 2 * time.Second maxBackoff := 30 * time.Second + var err error for attempt := 1; attempt <= maxRetries; attempt++ { - nodeID, apiToken, natsJWT, natsSeed, err = c.Register(ctx, body) + var res *RegisterResponse + res, err = c.RegisterFull(ctx, body) if err == nil { - return nodeID, apiToken, natsJWT, natsSeed, nil + return res, nil } if attempt == maxRetries { - return "", "", "", "", fmt.Errorf("failed after %d attempts: %w", maxRetries, err) + return nil, fmt.Errorf("failed after %d attempts: %w", maxRetries, err) } xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) select { case <-ctx.Done(): - return "", "", "", "", ctx.Err() + return nil, ctx.Err() case <-time.After(backoff): } backoff = min(backoff*2, maxBackoff) } - return nodeID, apiToken, natsJWT, natsSeed, err + return nil, err } // Heartbeat sends a single heartbeat POST with the given body. diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 24dd6f3c8ed7..408db9c88b95 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -15,6 +15,11 @@ import ( // package (and its gorm/DB dependencies). const statusPending = "pending" +// carrierTunnel mirrors cluster.CarrierTunnel, for the reason statusPending is +// copied: this package does not import the cluster package and its database +// dependencies. +const carrierTunnel = "tunnel" + // defaultMaxAttempts bounds how many times Acquire registers (and how many // consecutive times RefreshLoop may fail) before giving up. It is high enough // to ride out a slow admin approval or a transient frontend outage, but finite @@ -25,20 +30,21 @@ const defaultMaxAttempts = 100 // RegisterFunc performs one idempotent registration round-trip. type RegisterFunc func(ctx context.Context) (*RegisterResponse, error) -// NATSCredentialManager acquires NATS credentials at startup — waiting through -// admin approval when required — and refreshes them before the minted JWT -// expires, by re-registering (which mints a fresh JWT). The live NATS -// connection adopts a refreshed JWT on its next reconnect via Provider. Safe -// for concurrent use. +// CredentialManager acquires the credentials of the carrier at startup. For +// NATS it waits through admin approval when required, and it refreshes the +// credentials before the minted JWT expires, by registering again (which mints a +// new JWT). The live NATS connection adopts a refreshed JWT on its next +// reconnect via Provider. For the tunnel it keeps the token of the node that the +// last registration returned. Safe for concurrent use. // // It addresses two failure modes: a worker that needs credentials but registers // while still pending approval (it would otherwise give up and never connect), // and a long-running worker whose 24h JWT expires with no way to renew it. -type NATSCredentialManager struct { +type CredentialManager struct { register RegisterFunc requireCreds bool // block until credentials are present (frontend minting in use) - // Tunables; defaults set by NewNATSCredentialManager, overridable in tests. + // Tunables; defaults set by NewCredentialManager, overridable in tests. initialBackoff time.Duration maxBackoff time.Duration maxAttempts int // bound on Acquire attempts / consecutive refresh failures (<=0 = unlimited) @@ -46,16 +52,18 @@ type NATSCredentialManager struct { refreshRetry time.Duration expiryOf func(jwt string) (time.Time, bool) - mu sync.RWMutex - jwt string - seed string - nodeID string + mu sync.RWMutex + jwt string + seed string + nodeID string + carrier string + tunnelToken string } -// NewNATSCredentialManager builds a manager over register. When requireCreds is +// NewCredentialManager builds a manager over register. When requireCreds is // true, Acquire blocks until the node is approved and credentials are minted. -func NewNATSCredentialManager(register RegisterFunc, requireCreds bool) *NATSCredentialManager { - return &NATSCredentialManager{ +func NewCredentialManager(register RegisterFunc, requireCreds bool) *CredentialManager { + return &CredentialManager{ register: register, requireCreds: requireCreds, initialBackoff: 2 * time.Second, @@ -80,24 +88,45 @@ func jwtExpiry(token string) (time.Time, bool) { return time.Unix(uc.Expires, 0), true } -func (m *NATSCredentialManager) store(res *RegisterResponse) { +func (m *CredentialManager) store(res *RegisterResponse) { m.mu.Lock() defer m.mu.Unlock() m.nodeID = res.ID + m.carrier = res.Carrier if res.NatsJWT != "" && res.NatsUserSeed != "" { m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed } + if res.TunnelToken != "" { + m.tunnelToken = res.TunnelToken + } +} + +// Carrier returns the carrier that the last registration named. It is empty for +// a frontend that predates the field, and that is a frontend with NATS. +func (m *CredentialManager) Carrier() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.carrier +} + +// TunnelToken returns the token of the node for the tunnel, or an empty string +// before the first registration that returned one. The tunnel client calls it +// for every dial, because a registration mints a new token. +func (m *CredentialManager) TunnelToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.tunnelToken } // Current returns the latest NATS credentials (both empty until acquired). -func (m *NATSCredentialManager) Current() (jwt, seed string) { +func (m *CredentialManager) Current() (jwt, seed string) { m.mu.RLock() defer m.mu.RUnlock() return m.jwt, m.seed } // NodeID returns the node ID from the most recent registration. -func (m *NATSCredentialManager) NodeID() string { +func (m *CredentialManager) NodeID() string { m.mu.RLock() defer m.mu.RUnlock() return m.nodeID @@ -105,12 +134,12 @@ func (m *NATSCredentialManager) NodeID() string { // Provider returns a callback compatible with messaging.WithUserJWTProvider, // supplying the current credentials on each (re)connect. -func (m *NATSCredentialManager) Provider() func() (string, string) { +func (m *CredentialManager) Provider() func() (string, string) { return m.Current } // HasCredentials reports whether complete NATS credentials have been obtained. -func (m *NATSCredentialManager) HasCredentials() bool { +func (m *CredentialManager) HasCredentials() bool { jwt, seed := m.Current() return jwt != "" && seed != "" } @@ -119,7 +148,13 @@ func (m *NATSCredentialManager) HasCredentials() bool { // exponential backoff until the node is approved (status != pending) and // credentials are minted. Without requireCreds it returns the first successful // response (the historical one-shot behavior, preserved for anonymous NATS). -func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, error) { +// +// When the frontend names the tunnel as the carrier, the token for the tunnel is +// what Acquire waits for, and it does not matter whether the node is approved. +// The frontend mints the token at registration and the connect route refuses a +// node that waits for approval, so the tunnel client waits for the approval by +// itself. +func (m *CredentialManager) Acquire(ctx context.Context) (*RegisterResponse, error) { backoff := m.initialBackoff var lastReason error for attempt := 1; m.maxAttempts <= 0 || attempt <= m.maxAttempts; attempt++ { @@ -128,6 +163,12 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, case err != nil: lastReason = err xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) + case res.Carrier == carrierTunnel && res.TunnelToken == "": + lastReason = fmt.Errorf("node %s registered but the tunnel credential was not minted", res.ID) + xlog.Info("Node registered but the tunnel credential is not minted yet; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff) + case res.Carrier == carrierTunnel: + m.store(res) + return res, nil case !m.requireCreds: m.store(res) return res, nil @@ -158,7 +199,7 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, // error after maxAttempts consecutive refresh failures — letting the caller // exit the worker so it restarts and re-acquires (or surfaces the outage) // rather than silently drifting toward an expired, unrenewable JWT. -func (m *NATSCredentialManager) RefreshLoop(ctx context.Context) error { +func (m *CredentialManager) RefreshLoop(ctx context.Context) error { failures := 0 for { jwt, _ := m.Current() diff --git a/core/cli/workerregistry/credentials_test.go b/core/cli/workerregistry/credentials_test.go index 02ed413c4e8f..a1fd130a813c 100644 --- a/core/cli/workerregistry/credentials_test.go +++ b/core/cli/workerregistry/credentials_test.go @@ -50,7 +50,7 @@ func (f *fakeRegister) count() int { return f.calls } -var _ = Describe("NATSCredentialManager", func() { +var _ = Describe("CredentialManager", func() { approved := func(jwt, seed string) *RegisterResponse { return &RegisterResponse{ID: "node-1", Status: "healthy", NatsJWT: jwt, NatsUserSeed: seed} } @@ -63,7 +63,7 @@ var _ = Describe("NATSCredentialManager", func() { {res: approved("", "")}, // approved but JWT not minted yet {res: approved("jwt-1", "seed-1")}, // finally minted }} - m := NewNATSCredentialManager(f.fn(), true /* requireCreds */) + m := NewCredentialManager(f.fn(), true /* requireCreds */) m.initialBackoff = time.Millisecond m.maxBackoff = time.Millisecond @@ -81,7 +81,7 @@ var _ = Describe("NATSCredentialManager", func() { It("returns immediately on the first success when credentials are not required (anonymous NATS)", func() { f := &fakeRegister{steps: []step{{res: pending}}} - m := NewNATSCredentialManager(f.fn(), false /* requireCreds */) + m := NewCredentialManager(f.fn(), false /* requireCreds */) res, err := m.Acquire(context.Background()) Expect(err).ToNot(HaveOccurred()) @@ -92,7 +92,7 @@ var _ = Describe("NATSCredentialManager", func() { It("aborts when the context is cancelled while waiting for approval", func() { f := &fakeRegister{steps: []step{{res: pending}}} - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.initialBackoff = 10 * time.Millisecond ctx, cancel := context.WithCancel(context.Background()) @@ -103,7 +103,7 @@ var _ = Describe("NATSCredentialManager", func() { It("gives up after a bounded number of attempts so the worker exits and alerts", func() { f := &fakeRegister{steps: []step{{res: pending}}} // never approved - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.initialBackoff = time.Millisecond m.maxBackoff = time.Millisecond m.maxAttempts = 5 @@ -116,10 +116,72 @@ var _ = Describe("NATSCredentialManager", func() { }) }) + Describe("Acquire on the tunnel carrier", func() { + tunnelRes := func(status, token string) *RegisterResponse { + return &RegisterResponse{ID: "node-1", Status: status, Carrier: "tunnel", CarrierEpoch: 4, TunnelToken: token} + } + + It("returns at once with the token, also for a node that waits for approval", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("pending", "token-1")}}} + m := NewCredentialManager(f.fn(), true) + + res, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(res.Carrier).To(Equal("tunnel")) + Expect(f.count()).To(Equal(1)) + Expect(m.TunnelToken()).To(Equal("token-1")) + Expect(m.Carrier()).To(Equal("tunnel")) + Expect(m.HasCredentials()).To(BeFalse(), "a tunnel worker holds no NATS credentials") + }) + + It("waits when the frontend named the tunnel but did not mint a token", func() { + f := &fakeRegister{steps: []step{ + {res: tunnelRes("healthy", "")}, + {res: tunnelRes("healthy", "token-2")}, + }} + m := NewCredentialManager(f.fn(), false) + m.initialBackoff = time.Millisecond + m.maxBackoff = time.Millisecond + + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(f.count()).To(Equal(2)) + Expect(m.TunnelToken()).To(Equal("token-2")) + }) + + It("gives the token of the latest registration, because every registration mints a new one", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("healthy", "token-1")}, {res: tunnelRes("healthy", "token-2")}}} + m := NewCredentialManager(f.fn(), true) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(m.TunnelToken()).To(Equal("token-1")) + + m.store(&RegisterResponse{ID: "node-1", Carrier: "tunnel", TunnelToken: "token-2"}) + Expect(m.TunnelToken()).To(Equal("token-2")) + }) + + It("does not touch the NATS path: a frontend that names NATS, or nothing, still waits for the JWT", func() { + for _, carrier := range []string{"", "nats"} { + f := &fakeRegister{steps: []step{ + {res: &RegisterResponse{ID: "node-1", Status: "healthy", Carrier: carrier}}, + {res: &RegisterResponse{ID: "node-1", Status: "healthy", Carrier: carrier, NatsJWT: "jwt", NatsUserSeed: "seed"}}, + }} + m := NewCredentialManager(f.fn(), true) + m.initialBackoff = time.Millisecond + m.maxBackoff = time.Millisecond + + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(f.count()).To(Equal(2), "carrier %q", carrier) + Expect(m.TunnelToken()).To(BeEmpty()) + } + }) + }) + Describe("RefreshLoop (#5 — renew before the JWT expires)", func() { It("re-registers before expiry and updates the credentials served to new connections", func() { f := &fakeRegister{steps: []step{{res: approved("jwt-2", "seed-2")}}} - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.refreshLead = 0.5 m.refreshRetry = time.Millisecond // jwt-1 expires soon; jwt-2 is long-lived so the loop then idles. @@ -147,7 +209,7 @@ var _ = Describe("NATSCredentialManager", func() { It("returns an error after the bounded number of consecutive failures so the caller can exit", func() { f := &fakeRegister{steps: []step{{err: context.DeadlineExceeded}}} // refresh always fails - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.refreshLead = 0.5 m.refreshRetry = time.Millisecond m.maxAttempts = 3 @@ -161,7 +223,7 @@ var _ = Describe("NATSCredentialManager", func() { It("exits promptly when the current credential has no expiry (nothing to refresh)", func() { f := &fakeRegister{steps: []step{{res: approved("x", "y")}}} - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.expiryOf = func(string) (time.Time, bool) { return time.Time{}, false } m.store(approved("static", "seed")) From c0370924b5311e1a86e162ac4114b6284fd20164 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:50:29 +0000 Subject: [PATCH 027/162] feat(worker): add the client of the worker tunnel The worker dials the frontend and serves the streams that the frontend opens. It holds two sessions. The inference lane carries model calls and health. The bulk lane carries file transfers and lives only as long as the inference lane. A dial of the bulk lane that lands on a replica which does not hold the inference lane gets 409, and the worker tries again soon. A frontend that refuses the bulk lane does not cost the worker its inference lane. A stream is routed by the tag in its first frame. The host that the frontend names is dropped, and the port must be in the range of the allocator of the worker, so the tunnel cannot reach the LAN of the worker. A refusal says which of four things is true, and a late frame, a timeout and a cancelled context are never evidence about a backend. The wait between two attempts grows from half a second to thirty seconds with jitter, and it goes back to the floor only after a session that lasted. The token is read at every dial. The specs run the client against the real connect endpoint and a real PostgreSQL. Behind a link of 1 Gbit/s with 10 ms of delay in each direction, a transfer of 256 MiB delays a small call on the same session to a p99 of 88 ms, and on the bulk lane to 34 ms, against 31 ms for a call on its own TCP connection. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/tunnel.go | 859 ++++++++++++++++ core/services/worker/tunnel_endpoint_test.go | 372 +++++++ core/services/worker/tunnel_size_race_test.go | 5 + core/services/worker/tunnel_size_test.go | 8 + core/services/worker/tunnel_test.go | 946 ++++++++++++++++++ 5 files changed, 2190 insertions(+) create mode 100644 core/services/worker/tunnel.go create mode 100644 core/services/worker/tunnel_endpoint_test.go create mode 100644 core/services/worker/tunnel_size_race_test.go create mode 100644 core/services/worker/tunnel_size_test.go create mode 100644 core/services/worker/tunnel_test.go diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go new file mode 100644 index 000000000000..99652e1a3dd1 --- /dev/null +++ b/core/services/worker/tunnel.go @@ -0,0 +1,859 @@ +package worker + +import ( + "cmp" + "context" + "errors" + "fmt" + "math/rand/v2" + "net" + "net/http" + "net/url" + "os" + "strconv" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/tunnel" +) + +// The worker end of the tunnel. +// +// The worker dials out and never listens. It holds a websocket to a frontend and +// a multiplexed session on it, and every request that the frontend makes of this +// worker arrives as a stream in that session. A worker behind NAT, in another +// cluster or on a laptop needs no inbound port and no address that others can +// reach. +// +// A worker holds two sessions. The inference lane carries model calls, control +// and health. The bulk lane carries file transfers, so that a large transfer +// does not delay a small call. The bulk lane exists only while the inference +// lane exists: the frontend accepts it only on the replica that holds the +// inference lane, and the worker closes it with the inference lane. +// +// This side is the yamux client and it only accepts streams. The frontend is the +// server and it only opens them. The frontend asks and the worker answers. + +const ( + // tunnelBackoffBase is the shortest wait between two attempts to connect, + // before the jitter. + tunnelBackoffBase = 500 * time.Millisecond + + // tunnelBackoffMax is the longest wait. Without a limit, a worker that sits + // through a long outage of the frontend waits for hours and does not come + // back for a long time after the frontend does. The jitter and the floor stop + // a fleet of workers from turning a rolling restart into a storm of retries + // against the first replica that is back. + tunnelBackoffMax = 30 * time.Second + + // tunnelHealthyAfter is how long a session must last before the wait goes + // back to its floor. If the wait went back at once on connect, a replica that + // accepts the dial and dies a moment later would give a worker that retries + // at the floor for ever. The value is the keepalive interval of yamux, which + // is the shortest time over which a session that is only up can be told from + // one that works. + tunnelHealthyAfter = 30 * time.Second + + // tunnelHandshakeTimeout bounds the websocket upgrade. + tunnelHandshakeTimeout = 10 * time.Second + + // tunnelWrongReplicaDelay is the base of the wait after the answer 409 on the + // bulk lane. The answer means that the load balancer sent the dial to a + // replica that does not hold the inference lane. Another dial can land on the + // right one, so the wait is short. + tunnelWrongReplicaDelay = 250 * time.Millisecond + + // tunnelHeaderTimeout bounds how long a stream may take to send its request + // frame. Without a bound, a stream that sends nothing would hold a goroutine + // and a slot of the session for as long as the tunnel lives. + // + // The bound is long on purpose. On the relay path the timer starts when the + // owning replica opens the stream, and the dialling replica writes the frame + // only after the acceptance of the relay has come back to it. A round trip + // of a link between replicas, which carries artifacts of many gigabytes + // beside token streams, is inside this time. An expiry says nothing about + // the frontend, so it is refused with tunnel.ErrStreamNotServed. See accept. + tunnelHeaderTimeout = 15 * time.Second +) + +// LocalService opens a connection to a service that runs on this worker. +// +// target is the argument of the tag in the request frame of the stream. The +// service decides what it accepts: a frontend that names an address does not +// oblige the worker to dial it. See loopbackService. +type LocalService func(ctx context.Context, target string) (net.Conn, error) + +// TunnelConfig configures the tunnel that a worker holds to a frontend. +type TunnelConfig struct { + // FrontendURL is the URL that the worker registers with. The tunnel changes + // its scheme to ws or wss. + FrontendURL string + + // NodeID is the identity that registration gave to this worker. + NodeID string + + // Token returns the own tunnel credential of the node. The credential is new + // after every registration, so a client that kept the value from its start + // would present a value that the frontend no longer accepts, as soon as + // anything registered this worker again. Token is called once for every dial. + Token func() string + + // Services routes a stream by the tag in its request frame. A tag with no + // entry is refused. See accept. + Services map[string]LocalService + + // The fields below are for the specs. They are not exported, so that a caller + // cannot reach them. + sleep func(ctx context.Context, d time.Duration) error + now func() time.Time + headerTimeout time.Duration + // noBulk stops the worker from dialling the bulk lane. A spec that measures + // the waits of the inference lane needs the waits of the other lane out of + // its way. + noBulk bool +} + +// Tunnel is a running worker tunnel. One goroutine holds the inference session +// and connects it again when it ends, until Close. +type Tunnel struct { + endpoint string + nodeID string + token func() string + services map[string]LocalService + dialer *websocket.Dialer + + headerTimeout time.Duration + noBulk bool + sleep func(ctx context.Context, d time.Duration) error + // now measures how long a session lasted and nothing else. Deadlines use + // time.Now. A spec that fakes this clock to test the wait must not move + // every deadline of the package. + now func() time.Time + + cancel context.CancelFunc + done chan struct{} + closeOnce sync.Once + + // mu guards the sessions. The loops write them. Connected and + // BulkConnected read them, and on a running worker the reader is the + // handler of /readyz. + mu sync.Mutex + inference *tunnel.Session + bulk *tunnel.Session +} + +// Connected reports whether the tunnel holds an inference session that has not +// ended. +// +// It is false between two sessions and while the first dial is in progress. That +// is a statement about reachability and nothing else. A worker that connects +// again after a restart of its frontend reports false and is still a registered, +// running node. Nothing may read it as the worker being gone. +func (t *Tunnel) Connected() bool { + if t == nil { + return false + } + t.mu.Lock() + sess := t.inference + t.mu.Unlock() + // A session that ended is the value of the field until the loop clears it. + // The gap between the two is the time in which a probe must not answer 200. + return sess != nil && !sess.IsClosed() +} + +// BulkConnected reports whether the tunnel holds a bulk session that has not +// ended. +func (t *Tunnel) BulkConnected() bool { + if t == nil { + return false + } + t.mu.Lock() + sess := t.bulk + t.mu.Unlock() + return sess != nil && !sess.IsClosed() +} + +func (t *Tunnel) setSession(lane tunnel.Lane, sess *tunnel.Session) { + t.mu.Lock() + defer t.mu.Unlock() + if lane == tunnel.LaneBulk { + t.bulk = sess + return + } + t.inference = sess +} + +// StartTunnel dials the frontend and holds the tunnel until ctx ends or Close is +// called. +// +// The error is about the configuration and never about the frontend. A frontend +// that is down, that has not been upgraded or that refuses the credential is no +// reason for a worker not to start: it tries again in the background, with a +// wait that grows. If a failed dial made the start fail, the restart of a +// frontend would become an outage of the whole fleet of workers. +func StartTunnel(ctx context.Context, cfg TunnelConfig) (*Tunnel, error) { + if cfg.NodeID == "" { + return nil, errors.New("starting the worker tunnel: no node id") + } + if cfg.Token == nil { + return nil, errors.New("starting the worker tunnel: no credential source") + } + endpoint, err := tunnelEndpoint(cfg.FrontendURL, cfg.NodeID) + if err != nil { + return nil, err + } + + // Copied, so that the routing table cannot change under the accept loop. + services := make(map[string]LocalService, len(cfg.Services)) + for tag, svc := range cfg.Services { + services[tag] = svc + } + + dialer := tunnel.NewDialer(tunnelHandshakeTimeout) + // A worker reaches its frontend over the internet in the deployments that + // this is for, so it honours the proxy settings of the environment. + dialer.Proxy = http.ProxyFromEnvironment + + t := &Tunnel{ + endpoint: endpoint, + nodeID: cfg.NodeID, + token: cfg.Token, + services: services, + dialer: dialer, + headerTimeout: cmp.Or(cfg.headerTimeout, tunnelHeaderTimeout), + noBulk: cfg.noBulk, + sleep: cfg.sleep, + now: cfg.now, + done: make(chan struct{}), + } + if t.sleep == nil { + t.sleep = tunnelSleep + } + if t.now == nil { + t.now = time.Now + } + + loopCtx, cancel := context.WithCancel(ctx) + t.cancel = cancel + go func() { + defer close(t.done) + t.run(loopCtx) + }() + return t, nil +} + +// Close stops the tunnel and waits for its loop to end. It is idempotent. +func (t *Tunnel) Close() error { + t.closeOnce.Do(func() { + t.cancel() + <-t.done + }) + return nil +} + +// run holds one inference session at a time and connects again with a wait that +// is bounded. +func (t *Tunnel) run(ctx context.Context) { + attempt := 0 + for { + if ctx.Err() != nil { + return + } + + start := t.now() + err := t.connectAndServe(ctx) + if ctx.Err() != nil { + return + } + + // A session that lasted is the only sign that the frontend is well. See + // tunnelHealthyAfter for why a connect is not. + if t.now().Sub(start) >= tunnelHealthyAfter { + attempt = 0 + } + attempt++ + + delay := tunnelBackoffDelay(attempt) + t.logSessionEnded(err, attempt, delay) + if err := t.sleep(ctx, delay); err != nil { + return + } + } +} + +// connectAndServe dials the inference lane, serves streams until the session +// ends, and leaves nothing running behind it. +func (t *Tunnel) connectAndServe(ctx context.Context) error { + sess, err := t.dial(ctx, tunnel.LaneInference) + if err != nil { + return err + } + xlog.Info("Worker tunnel established", "node", t.nodeID, "frontend", t.endpoint) + + // Published before the accept loop starts. Connected also checks IsClosed, + // and that check is what keeps the answer right after the session ended: the + // clear below runs after the streams that are in progress have ended, and the + // probe must say "not ready" during that wait. + t.setSession(tunnel.LaneInference, sess) + defer t.setSession(tunnel.LaneInference, nil) + + // Streams run under the context of the session and not under the one of the + // loop. A stream that waits for a local dial would outlive the session and + // hold the next connect behind it. + sessCtx, endSession := context.WithCancel(ctx) + + // AcceptStream takes no context, so closing the session is what stops it + // when the worker shuts down. + watchdogDone := make(chan struct{}) + go func() { + defer close(watchdogDone) + select { + case <-sessCtx.Done(): + _ = sess.Close() + case <-sess.CloseChan(): + } + }() + + var streams sync.WaitGroup + // The bulk lane lives as long as this inference session. + if !t.noBulk { + streams.Go(func() { t.runBulk(sessCtx, &streams) }) + } + + serveErr := t.serve(sessCtx, sess, &streams) + + endSession() + _ = sess.Close() + <-watchdogDone + // Closing the session ends every stream goroutine: yamux closes the streams + // of a session that closes, which wakes a read and fails a write. The wait + // keeps a new connect from overlapping the streams of the session it + // replaced. + streams.Wait() + return serveErr +} + +// runBulk holds the bulk session for as long as ctx lives. It dials again when +// the session ends, and it does not stop the inference lane when it fails. A +// frontend that predates the bulk lane answers a dial to it as an inference +// dial, or refuses it. Either way the worker must still serve. +func (t *Tunnel) runBulk(ctx context.Context, streams *sync.WaitGroup) { + attempt := 0 + for ctx.Err() == nil { + start := t.now() + sess, err := t.dial(ctx, tunnel.LaneBulk) + var delay time.Duration + switch { + case err == nil: + xlog.Info("Worker bulk tunnel established", "node", t.nodeID) + t.setSession(tunnel.LaneBulk, sess) + serveErr := t.serveUntilDone(ctx, sess, streams) + t.setSession(tunnel.LaneBulk, nil) + if ctx.Err() != nil { + return + } + if t.now().Sub(start) >= tunnelHealthyAfter { + attempt = 0 + } + attempt++ + delay = tunnelBackoffDelay(attempt) + xlog.Debug("worker bulk tunnel ended", "node", t.nodeID, "retry_in", delay, "error", serveErr) + case isWrongReplica(err): + // The dial landed on a replica that does not hold the inference + // lane. Not a failure, and not an attempt for the wait. + delay = wrongReplicaDelay() + xlog.Debug("worker bulk tunnel dial reached a replica that does not hold the tunnel", "node", t.nodeID, "retry_in", delay) + default: + if ctx.Err() != nil { + // The inference session ended during the dial. That is not a + // failure of the bulk lane. + return + } + attempt++ + delay = tunnelBackoffDelay(attempt) + xlog.Debug("worker bulk tunnel dial failed", "node", t.nodeID, "retry_in", delay, "error", err) + } + if err := t.sleep(ctx, delay); err != nil { + return + } + } +} + +// serveUntilDone serves the streams of a session until it ends or ctx ends, and +// returns after the streams have ended. +func (t *Tunnel) serveUntilDone(ctx context.Context, sess *tunnel.Session, streams *sync.WaitGroup) error { + done := make(chan struct{}) + defer close(done) + go func() { + select { + case <-ctx.Done(): + _ = sess.Close() + case <-done: + } + }() + err := t.serve(ctx, sess, streams) + _ = sess.Close() + return err +} + +// serve accepts streams until the session ends. +// +// A goroutine serves each stream, and the error of a stream never reaches this +// loop. A request that is malformed or cannot be routed must not cost the worker +// the other requests of the session. +func (t *Tunnel) serve(ctx context.Context, sess *tunnel.Session, streams *sync.WaitGroup) error { + for { + stream, err := sess.AcceptStream() + if err != nil { + return err + } + streams.Go(func() { t.handleStream(ctx, stream) }) + } +} + +// handleStream reads the request frame of a stream and then joins the stream to +// a local service, or refuses it. +func (t *Tunnel) handleStream(ctx context.Context, stream net.Conn) { + // A panic under one stream must not take the worker down. Nothing supervises + // this goroutine, so a panic that is not recovered ends the process, and with + // it the session and every other stream. This does not panic again, because + // no middleware above it would report the panic. It covers what runs in this + // goroutine: the request frame, the lookup of the route, and the dial of the + // local service. It does not cover the copy goroutines of Splice. + defer func() { + if r := recover(); r != nil { + xlog.Error("Panic while serving a worker tunnel stream", "node", t.nodeID, "panic", r) + _ = stream.Close() + } + }() + + local, ok := t.accept(ctx, stream) + if !ok { + // accept has answered and closed the stream. + return + } + + // Splice closes both ends. + if err := tunnel.Splice(stream, local); err != nil { + xlog.Debug("worker tunnel stream ended with an error", "node", t.nodeID, "error", err) + } +} + +// accept reads the request frame and resolves it to a local connection. The +// result is false when the stream was refused: the refusal was sent and the +// stream closed. +func (t *Tunnel) accept(ctx context.Context, stream net.Conn) (net.Conn, bool) { + if err := stream.SetReadDeadline(time.Now().Add(t.headerTimeout)); err != nil { + // Not served, and not unavailable: this is a fact about the stream, which + // took no deadline. No local service was named yet, so there is nothing + // to call unreachable. + t.refuse(stream, fmt.Errorf("%w: arming the request deadline: %v", tunnel.ErrStreamNotServed, err)) + return nil, false + } + + tag, target, err := tunnel.ReadStreamRequest(stream) + if err != nil { + // Two causes, and they stay apart. A frame that is malformed is a bug of + // the frontend and does not clear, so it is evidence that the frontend + // acts on. A deadline that expired is a frame that has not arrived yet, + // and that clears when the link has room. A model that is evicted + // across the fleet because a link was congested is the failure that this + // distinction prevents. + if reportsTimeout(err) { + t.refuse(stream, fmt.Errorf("%w: %v", tunnel.ErrStreamNotServed, err)) + return nil, false + } + t.refuse(stream, fmt.Errorf("%w: %v", tunnel.ErrStreamRequestInvalid, err)) + return nil, false + } + + svc, known := t.services[tag] + if !known { + // A fact about what this worker serves. The frontend knows that a retry + // is useless until the worker is upgraded, which is not what it should + // conclude from "unavailable". + t.refuse(stream, fmt.Errorf("%w: %q", tunnel.ErrStreamTagUnknown, tag)) + return nil, false + } + + // The deadline is cleared before the local dial. What follows the request + // frame belongs to the carried protocol, which has its own deadlines, and a + // deadline that stayed would end a long inference stream. + if err := stream.SetReadDeadline(time.Time{}); err != nil { + t.refuse(stream, fmt.Errorf("%w: clearing the request deadline: %v", tunnel.ErrStreamNotServed, err)) + return nil, false + } + + local, err := svc(ctx, target) + if err != nil { + t.refuse(stream, classifyServiceFailure(err)) + return nil, false + } + + if err := tunnel.WriteStreamAccepted(stream); err != nil { + // The frontend never learns that the stream was accepted, so it cannot be + // used. The local connection is closed here, or it would leak once for + // every reply that failed. + xlog.Debug("worker tunnel could not accept a stream", "node", t.nodeID, "error", err) + _ = local.Close() + _ = stream.Close() + return nil, false + } + return local, true +} + +// classifyServiceFailure decides which refusal the error of a local service is. +// +// A service that classified its failure keeps its classification, and that is +// asked of the whole vocabulary. loopbackService classifies, and the difference +// matters: a target outside the port range of the worker is a request that this +// worker never serves, and a backend that does not listen yet is a condition +// that clears. The first reported as the second makes the frontend retry what +// cannot work. The second reported as the first makes it give up on a backend +// that is starting. +// +// The default is "unavailable", and that is on purpose, because the frontend +// acts on that code. A dial to the process that failed is the nearest thing to +// evidence that the worker can produce. The direction of a wrong classification +// decides which mistake is made. A wrong eviction is active: it deletes rows +// and shuts down a model that is loaded. A wrong retention is passive, is +// limited to one slot of one replica, and clears at a restart or an eviction. +// Because of this, the exceptions below are a deny list of causes that are not +// the target answering, and not an allow list of causes that are. A deny list +// that misses a cause fails loudly. An allow list that misses one fails +// silently, for ever. +// +// The causes that are exempt: +// +// - The end of the context. It is the context of the session, and it ends +// while stream goroutines are running, so this worker is closing its +// tunnel and will connect again. +// - The own deadline of this process (os.ErrDeadlineExceeded). +// - Anything else that says of itself that it is a timeout. On a dial this +// includes ETIMEDOUT and EAGAIN. EAGAIN is this worker running out of a +// resource, and that is not about the target. ETIMEDOUT from connect(2) is +// an observation of the target, but it is "the handshake did not finish", +// which is a listener that is blocked or full and not one that is absent. +// An eviction of a backend that is overloaded is what this design must +// prevent. ECONNREFUSED, the shape of a process that is gone, is not a +// timeout and still gives "unavailable". +// +// It must never give the unknown-tag refusal: a tag that the worker serves does +// not stop being served because one dial failed. +func classifyServiceFailure(err error) error { + if tunnel.IsStreamRefusal(err) { + return err + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || reportsTimeout(err) { + return fmt.Errorf("%w: %v", tunnel.ErrStreamNotServed, err) + } + return fmt.Errorf("%w: %v", tunnel.ErrStreamTargetUnavailable, err) +} + +// reportsTimeout reports whether err says of itself that it is a timeout. +// +// It is asked in two places that mean different things. On a read of a stream it +// is a deadline that this process set. On a dial it is wider: Go reports ETIMEDOUT +// and EAGAIN as timeouts, and net.OpError passes that on. Both places want the +// same answer, "not the target speaking", so one predicate serves both. See +// classifyServiceFailure. +// +// It asks net.Error as well as os.ErrDeadlineExceeded because the two sets +// differ: a yamux stream returns its own timeout value from a read whose +// deadline passed, and a net.OpError on a socket returns os.ErrDeadlineExceeded. +func reportsTimeout(err error) bool { + if errors.Is(err, os.ErrDeadlineExceeded) { + return true + } + var netErr net.Error + return errors.As(err, &netErr) && netErr.Timeout() +} + +// refuse says why a stream is not served and then ends it. +// +// The close is not optional. A worker that says why and leaves the stream open +// has left the frontend with a request that is never answered, which looks like +// a worker that is slow, and a deadline on the other side cannot tell the two +// apart. The reply makes the refusal readable and the close makes it prompt. The +// reply is best effort and the close is not. +func (t *Tunnel) refuse(stream net.Conn, reason error) { + if err := tunnel.WriteStreamRefusal(stream, reason); err != nil { + xlog.Debug("worker tunnel could not report why it refused a stream", "node", t.nodeID, "error", err) + } + _ = stream.Close() + xlog.Debug("worker tunnel refused a stream", "node", t.nodeID, "reason", reason) +} + +// dial opens the websocket for a lane and starts the session on it. +func (t *Tunnel) dial(ctx context.Context, lane tunnel.Lane) (*tunnel.Session, error) { + // Read here, once for every dial. See TunnelConfig.Token. + token := t.token() + if token == "" { + // Not a dial that fails with "unauthorized". The worker has no + // credential yet, and an operator who reads "unauthorized" looks for a + // mismatch that does not exist. + return nil, errors.New("dialling the worker tunnel: this node has no tunnel credential yet") + } + header := http.Header{} + header.Set("Authorization", "Bearer "+token) + + endpoint := t.endpoint + if lane == tunnel.LaneBulk { + endpoint += "&lane=" + url.QueryEscape(string(lane)) + } + ws, resp, err := t.dialer.DialContext(ctx, endpoint, header) + if err != nil { + if resp != nil { + // gorilla gives ErrBadHandshake for every status that is not 101. + // Without the status, a 401, a 403 and a 503 are one log line. + defer func() { _ = resp.Body.Close() }() + return nil, &tunnelDialError{status: resp.StatusCode, cause: err} + } + return nil, fmt.Errorf("dialling the worker tunnel: %w", err) + } + sess, err := tunnel.ClientSession(ws, lane) + if err != nil { + _ = ws.Close() + return nil, err + } + return sess, nil +} + +// tunnelDialError carries the status of a dial that the frontend refused, so +// that the refusals are not one line in the log. +type tunnelDialError struct { + status int + cause error +} + +func (e *tunnelDialError) Error() string { + return fmt.Sprintf("dialling the worker tunnel: frontend answered %d: %v", e.status, e.cause) +} + +func (e *tunnelDialError) Unwrap() error { return e.cause } + +// isWrongReplica reports whether the frontend answered 409 to a dial. +func isWrongReplica(err error) bool { + var dialErr *tunnelDialError + return errors.As(err, &dialErr) && dialErr.status == http.StatusConflict +} + +// wrongReplicaDelay returns a short wait with jitter. +func wrongReplicaDelay() time.Duration { + // #nosec G404 -- the jitter spreads dials. An attacker gains nothing by + // predicting it. + return tunnelWrongReplicaDelay/2 + time.Duration(rand.Int64N(int64(tunnelWrongReplicaDelay))) +} + +// logSessionEnded says why the tunnel connects again, at a level that fits what +// the operator can do about it. +// +// "Awaiting approval", "your token is wrong" and "this frontend has no tunnels" +// send an operator to three places, and a worker retries all three in the same +// way. None of them is a reason to stop, because a registration or an action of +// an admin fixes each without a restart of the worker. +func (t *Tunnel) logSessionEnded(err error, attempt int, delay time.Duration) { + if err == nil { + xlog.Info("Worker tunnel closed, connecting again", "node", t.nodeID, "attempt", attempt, "retry_in", delay) + return + } + + var dialErr *tunnelDialError + if errors.As(err, &dialErr) { + switch dialErr.status { + case http.StatusUnauthorized: + // Two causes, and this worker cannot fix either alone. It has no + // inbound listener, so a tunnel that it cannot open is a worker that + // nothing can reach. That is an outage and not a warning about a + // path that is worse. + xlog.Warn("Frontend rejected this worker's tunnel credential, so nothing can reach this worker; "+ + "either another worker registered under this node name and rotated the credential (check LOCALAI_NODE_NAME is unique), "+ + "or the frontend's record of this node was replaced. Restarting this worker registers it again and mints a new credential", + "node", t.nodeID, "retry_in", delay) + case http.StatusForbidden: + xlog.Info("Worker tunnel refused: this node is awaiting admin approval", + "node", t.nodeID, "retry_in", delay) + case http.StatusNotFound: + xlog.Debug("frontend does not serve worker tunnels, so it predates them", + "node", t.nodeID, "retry_in", delay) + case http.StatusServiceUnavailable: + xlog.Debug("frontend is not running in distributed mode, so it holds no worker tunnels", + "node", t.nodeID, "retry_in", delay) + default: + xlog.Warn("Worker tunnel dial refused", "node", t.nodeID, "status", dialErr.status, + "attempt", attempt, "retry_in", delay, "error", err) + } + return + } + xlog.Warn("Worker tunnel ended, connecting again", "node", t.nodeID, "attempt", attempt, "retry_in", delay, "error", err) +} + +// tunnelBackoffDelay returns the wait before attempt n. +// +// Half of the wait is fixed and half is drawn. Full jitter draws over the whole +// interval and can give a wait near zero, and a worker that can draw that can +// spin. The fixed half is a floor for every worker, and the drawn half keeps a +// fleet that lost the same replica from connecting again at the same moment. +func tunnelBackoffDelay(attempt int) time.Duration { + if attempt < 1 { + attempt = 1 + } + d := tunnelBackoffMax + // Two guards. The bound on attempt keeps the shift defined. The check for a + // positive value catches the overflow that would give a long outage a + // negative wait, which is a tight loop that looks like a backoff. + if attempt <= 40 { + if scaled := tunnelBackoffBase << (attempt - 1); scaled > 0 && scaled < tunnelBackoffMax { + d = scaled + } + } + // #nosec G404 -- the jitter spreads attempts so that a fleet does not retry + // together. An attacker gains nothing by predicting it. + return d/2 + time.Duration(rand.Int64N(int64(d/2)+1)) +} + +// tunnelSleep waits for d, or returns early when ctx ends. +func tunnelSleep(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// tunnelEndpoint turns the URL that a worker registers with into the websocket +// URL of its tunnel. The URL ends with a query, so a lane is added with "&". +func tunnelEndpoint(frontendURL, nodeID string) (string, error) { + if frontendURL == "" { + return "", errors.New("starting the worker tunnel: no frontend URL") + } + u, err := url.Parse(frontendURL) + if err != nil { + return "", fmt.Errorf("starting the worker tunnel: parsing frontend URL %q: %w", frontendURL, err) + } + switch u.Scheme { + case "http", "ws": + u.Scheme = "ws" + case "https", "wss": + u.Scheme = "wss" + default: + return "", fmt.Errorf("starting the worker tunnel: frontend URL %q has scheme %q, want http or https", frontendURL, u.Scheme) + } + if u.Host == "" { + return "", fmt.Errorf("starting the worker tunnel: frontend URL %q has no host", frontendURL) + } + // Appended and not assigned, so that a frontend behind a path prefix keeps + // it. Registration builds its URLs in the same way. + u.Path = strings.TrimRight(u.Path, "/") + tunnel.ConnectPath + u.RawQuery = url.Values{"id": []string{nodeID}}.Encode() + return u.String(), nil +} + +// loopbackService routes a tagged stream to a process that listens on the +// loopback interface of this worker. +// +// The host that the frontend names is discarded and only the port is used. This +// is the security property of the function. A tunnel ends inside the worker +// process, so a stream that arrives on it can reach whatever the worker can +// reach. Without this, whoever holds the frontend end could make every worker of +// the fleet dial any host of its private network, and the tunnel would be a +// proxy into the LAN of the worker. Without the host, the set that can be +// reached is this machine. +// +// The port range is the one that the port allocator of the worker gives to its +// backend processes, so a stream cannot be pointed at another service that +// happens to listen on this host. It is as narrow as that range. By default the +// range goes to 65535, and a deployment that wants it narrow sets +// LOCALAI_GRPC_MAX_PORT, which narrows both. +// +// The shape matters as much as the checks. Nothing that comes from the wire +// reaches the dialler. The address is built from the constant loopbackHost and +// from strconv.Itoa of an int that this function checked, so target has no path +// to DialContext. To turn this into a dialler for any host, a data flow has to +// be added, and a deleted check is not enough. That is the difference between a +// property and a guard. +func loopbackService(minPort, maxPort int) LocalService { + return func(ctx context.Context, target string) (net.Conn, error) { + _, portStr, err := net.SplitHostPort(target) + if err != nil { + return nil, fmt.Errorf("%w: routing a tunnel stream: %q is not a host:port: %v", + tunnel.ErrStreamRequestInvalid, target, err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + return nil, fmt.Errorf("%w: routing a tunnel stream: %q has no numeric port: %v", + tunnel.ErrStreamRequestInvalid, target, err) + } + if port < minPort || port > maxPort { + // Invalid and not unavailable: no retry brings a port outside the + // range of the allocator into it. + return nil, fmt.Errorf("%w: routing a tunnel stream: port %d is outside this worker's backend range [%d, %d]", + tunnel.ErrStreamRequestInvalid, port, minPort, maxPort) + } + var d net.Dialer + return d.DialContext(ctx, "tcp", net.JoinHostPort(loopbackHost, strconv.Itoa(port))) + } +} + +// loopbackHost is the host on which every stream that the frontend can steer is +// dialled. It is a constant so that "a stream cannot choose where the worker +// dials" is a fact about the code and not a claim about its inputs. +// +// It is not the only host that this file dials. fixedService dials the address +// it was made with, and Run makes it from LOCALAI_HTTP_ADDR of this worker, +// which an operator may set to an address that can be routed. loopbackAddr +// replaces a wildcard bind and leaves an explicit host alone, because a server +// bound to one address is not reachable on another. So the http tag can dial a +// host that is not loopback. That host was set by the operator for the own +// server of this worker, and no stream names it, because fixedService ignores +// its target. What the design needs is that the frontend cannot steer the dial, +// and that holds for both tags. +const loopbackHost = "127.0.0.1" + +// tunnelServices builds the routing table that the worker installs on its +// tunnel. It is its own function so that the table can be tested: the table is +// the security boundary of the tunnel, and a table built inline in Run could be +// reached only by starting a worker. +func tunnelServices(cfg *Config, httpBindAddr string) map[string]LocalService { + basePort := cfg.effectiveBasePort() + return map[string]LocalService{ + // The frontend names a backend process by its port. The worker decides + // that only its own loopback is reachable, and only inside its own port + // range. + tunnel.StreamTagGRPC: loopbackService(basePort, cfg.effectiveMaxPort(basePort)), + tunnel.StreamTagHTTP: fixedService(loopbackAddr(httpBindAddr)), + } +} + +// fixedService routes a tagged stream to one address of this worker and ignores +// what the frontend named. A worker has one HTTP server and only the worker +// knows where it listens, so the frontend has nothing useful to say about the +// target and is not given the chance. +func fixedService(addr string) LocalService { + return func(ctx context.Context, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + } +} + +// loopbackAddr turns a bind address into one that reaches the same listener from +// inside this process. A server bound to 0.0.0.0 can be reached on loopback, but +// to dial 0.0.0.0 reaches it only by accident and not on every platform, so the +// wildcard is replaced and not dialled. +func loopbackAddr(bindAddr string) string { + host, port, err := net.SplitHostPort(bindAddr) + if err != nil { + return bindAddr + } + if host == "" || host == "0.0.0.0" || host == "::" || host == "[::]" { + return net.JoinHostPort(loopbackHost, port) + } + return bindAddr +} diff --git a/core/services/worker/tunnel_endpoint_test.go b/core/services/worker/tunnel_endpoint_test.go new file mode 100644 index 000000000000..08796481db18 --- /dev/null +++ b/core/services/worker/tunnel_endpoint_test.go @@ -0,0 +1,372 @@ +package worker + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "net" + "net/http" + "net/http/httptest" + "slices" + "sync" + "time" + + "github.com/labstack/echo/v4" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/tunnel/slowlink" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// These specs run the worker client against the real connect endpoint, the real +// registry of tunnels and a real PostgreSQL, over real sockets. +var _ = Describe("Worker tunnel against the connect endpoint", func() { + const ( + nodeToken = "own-tunnel-token" + linkRate = 125e6 // 1 Gbit/s + oneWay = 10 * time.Millisecond + ) + + var ( + ctx context.Context + cancel context.CancelFunc + clusterR *cluster.Registry + registry *tunnel.Registry + frontend *httptest.Server + nodeID string + + hijackedMu sync.Mutex + hijacked []net.Conn + ) + + BeforeEach(func() { + ctx, cancel = context.WithCancel(context.Background()) + DeferCleanup(cancel) + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + registry = tunnel.NewRegistry(clusterR, "replica-a") + + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + nodeID = node.ID + sum := sha256.Sum256([]byte(nodeToken)) + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hex.EncodeToString(sum[:]))).To(Succeed()) + + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(nodeReg, registry)) + frontend = httptest.NewUnstartedServer(e) + // The connections that were upgraded leave the server, so the server does + // not close them. The spec closes them to cut the tunnel. + frontend.Config.ConnState = func(c net.Conn, state http.ConnState) { + if state == http.StateHijacked { + hijackedMu.Lock() + hijacked = append(hijacked, c) + hijackedMu.Unlock() + } + } + frontend.Start() + DeferCleanup(frontend.Close) + }) + + startWorker := func(frontendURL string, services map[string]LocalService) *Tunnel { + GinkgoHelper() + t, err := StartTunnel(ctx, TunnelConfig{ + FrontendURL: frontendURL, + NodeID: nodeID, + Token: func() string { return nodeToken }, + Services: services, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(t.Close) + return t + } + + // tcpService starts a server and returns its address. + tcpService := func(handle func(net.Conn)) string { + GinkgoHelper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(lis.Close) + go func() { + for { + c, err := lis.Accept() + if err != nil { + return + } + go handle(c) + } + }() + return lis.Addr().String() + } + + echo8 := func(c net.Conn) { + defer func() { _ = c.Close() }() + _, _ = io.Copy(c, c) + } + // download sends transferSize bytes to whoever connects, and then closes. + download := func(c net.Conn) { + defer func() { _ = c.Close() }() + chunk := make([]byte, 256<<10) + for sent := 0; sent < transferBytes; sent += len(chunk) { + if _, err := c.Write(chunk); err != nil { + return + } + } + } + + It("holds both lanes on the replica that owns the node and carries streams on each", func() { + echoAddr := tcpService(echo8) + t := startWorker(frontend.URL, map[string]LocalService{ + tunnel.StreamTagGRPC: dialLocalTCP, + tunnel.StreamTagHTTP: dialLocalTCP, + }) + Eventually(t.Connected, "10s").Should(BeTrue()) + Eventually(t.BulkConnected, "10s").Should(BeTrue()) + Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + + for _, lane := range []tunnel.Lane{tunnel.LaneInference, tunnel.LaneBulk} { + st, err := registry.Open(ctx, nodeID, lane) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(st, tunnel.StreamTagHTTP, echoAddr)).To(Succeed()) + Expect(tunnel.ReadStreamReply(st)).To(Succeed()) + _, err = st.Write([]byte("hello")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 5) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred(), string(lane)) + Expect(string(got)).To(Equal("hello")) + Expect(st.Close()).To(Succeed()) + } + }) + + It("connects both lanes again, with a new claim, after the connections to the frontend were cut", func() { + t := startWorker(frontend.URL, map[string]LocalService{}) + Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + Eventually(t.BulkConnected, "10s").Should(BeTrue()) + _, first, err := clusterR.Owner(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + + // The sockets are cut with no close frame, as they are when a replica is + // killed. + hijackedMu.Lock() + for _, c := range hijacked { + _ = c.Close() + } + hijackedMu.Unlock() + Eventually(t.Connected, "10s").Should(BeFalse()) + + Eventually(t.Connected, "30s").Should(BeTrue()) + Eventually(t.BulkConnected, "30s").Should(BeTrue()) + Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + _, second, err := clusterR.Owner(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(second).ToNot(Equal(first), "a new session is a new claim") + }) + + // The delay of a small call during a large transfer, on the real code: the + // connect endpoint, the registry, and the client of the worker, all behind a + // link with a rate and a delay. The numbers are the same kind as the ones in + // the spec of the tunnel package, and the lane is chosen by the caller of + // Open, as the file stager will do. + Describe("the delay of a small call during a transfer", Label("benchmark"), Ordered, func() { + const directName = "direct: transfer and probe on two TCP connections" + var results = map[string][]time.Duration{} + + // run starts a worker behind the link, and runs a download of + // transferBytes on transferLane while a probe runs on the inference lane. + run := func(name string, transferLane tunnel.Lane) { + GinkgoHelper() + echoAddr := tcpService(echo8) + downloadAddr := tcpService(download) + link, err := slowlink.New(frontend.Listener.Addr().String(), linkRate, oneWay) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(link.Close) + + t := startWorker("http://"+link.Addr(), map[string]LocalService{ + tunnel.StreamTagGRPC: dialLocalTCP, + tunnel.StreamTagHTTP: dialLocalTCP, + }) + Eventually(t.Connected, "10s").Should(BeTrue()) + Eventually(t.BulkConnected, "10s").Should(BeTrue()) + Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + + // The probe sends the request frame and the 8 bytes together, so a + // call costs one round trip, as a call on an open connection does. + probe := func() (time.Duration, error) { + start := time.Now() + st, err := registry.Open(ctx, nodeID, tunnel.LaneInference) + if err != nil { + return 0, err + } + defer func() { _ = st.Close() }() + if err := tunnel.WriteStreamRequest(st, tunnel.StreamTagGRPC, echoAddr); err != nil { + return 0, err + } + if _, err := st.Write(make([]byte, 8)); err != nil { + return 0, err + } + if err := tunnel.ReadStreamReply(st); err != nil { + return 0, err + } + if _, err := io.ReadFull(st, make([]byte, 8)); err != nil { + return 0, err + } + return time.Since(start), nil + } + + stop := make(chan struct{}) + sampled := make(chan []time.Duration, 1) + go func() { + var samples []time.Duration + for { + select { + case <-stop: + sampled <- samples + return + case <-time.After(5 * time.Millisecond): + } + d, err := probe() + if err != nil { + sampled <- samples + return + } + samples = append(samples, d) + } + }() + + start := time.Now() + st, err := registry.Open(ctx, nodeID, transferLane) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(st, tunnel.StreamTagHTTP, downloadAddr)).To(Succeed()) + Expect(tunnel.ReadStreamReply(st)).To(Succeed()) + n, err := io.Copy(io.Discard, st) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeEquivalentTo(transferBytes)) + elapsed := time.Since(start) + close(stop) + + var samples []time.Duration + Eventually(sampled, "30s").Should(Receive(&samples)) + Expect(len(samples)).To(BeNumerically(">=", 15), "too few probes completed during the transfer") + samples = samples[len(samples)/10:] + results[name] = samples + slices.Sort(samples) + ms := func(d time.Duration) float64 { return float64(d) / float64(time.Millisecond) } + AddReportEntry(name, map[string]any{ + "transfer MiB": transferBytes >> 20, + "transfer MiB/s": float64(transferBytes>>20) / elapsed.Seconds(), + "probes": len(samples), + "probe p50 ms": ms(samples[len(samples)/2]), + "probe p99 ms": ms(samples[min(len(samples)-1, len(samples)*99/100)]), + "probe max ms": ms(samples[len(samples)-1]), + }) + } + + It("measures the baseline: the transfer and the probe on two TCP connections", func() { + // One service behind the link. The first byte of a connection says + // what the connection is for. + front := tcpService(func(c net.Conn) { + kind := make([]byte, 1) + if _, err := io.ReadFull(c, kind); err != nil { + _ = c.Close() + return + } + if kind[0] == 'e' { + echo8(c) + } else { + download(c) + } + }) + link, err := slowlink.New(front, linkRate, oneWay) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(link.Close) + dial := func(kind byte) net.Conn { + c, err := net.Dial("tcp", link.Addr()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(c.Close) + _, err = c.Write([]byte{kind}) + Expect(err).ToNot(HaveOccurred()) + return c + } + bulk, probeConn := dial('d'), dial('e') + + stop := make(chan struct{}) + sampled := make(chan []time.Duration, 1) + go func() { + var samples []time.Duration + buf := make([]byte, 8) + for { + select { + case <-stop: + sampled <- samples + return + case <-time.After(5 * time.Millisecond): + } + start := time.Now() + if _, err := probeConn.Write(buf); err != nil { + sampled <- samples + return + } + if _, err := io.ReadFull(probeConn, buf); err != nil { + sampled <- samples + return + } + samples = append(samples, time.Since(start)) + } + }() + n, err := io.Copy(io.Discard, bulk) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeEquivalentTo(transferBytes)) + close(stop) + + var samples []time.Duration + Eventually(sampled, "30s").Should(Receive(&samples)) + Expect(len(samples)).To(BeNumerically(">=", 15)) + samples = samples[len(samples)/10:] + slices.Sort(samples) + results[directName] = samples + ms := func(d time.Duration) float64 { return float64(d) / float64(time.Millisecond) } + AddReportEntry(directName, map[string]any{ + "probes": len(samples), + "probe p50 ms": ms(samples[len(samples)/2]), + "probe p99 ms": ms(samples[min(len(samples)-1, len(samples)*99/100)]), + }) + }) + + It("measures the probe with the transfer on the same lane", func() { + run("shared: transfer and probe on the inference lane", tunnel.LaneInference) + }) + + It("measures the probe with the transfer on the bulk lane", func() { + run("bulk lane: transfer on the bulk lane, probe on the inference lane", tunnel.LaneBulk) + }) + + It("delays the probe less when the transfer uses the bulk lane", func() { + shared := results["shared: transfer and probe on the inference lane"] + bulk := results["bulk lane: transfer on the bulk lane, probe on the inference lane"] + direct := results[directName] + Expect(shared).ToNot(BeEmpty()) + Expect(bulk).ToNot(BeEmpty()) + Expect(direct).ToNot(BeEmpty()) + p99 := func(s []time.Duration) time.Duration { return s[min(len(s)-1, len(s)*99/100)] } + // The acceptance of the bulk lane: the probe stays within 10 ms of a + // probe on its own TCP connection, at the 99th percentile, behind + // the same link and the same queue. The same transfer on the lane of + // the probe delays it by tens of milliseconds more. + Expect(p99(bulk)).To(BeNumerically("<=", p99(direct)+10*time.Millisecond), + "bulk lane p99 %v, direct p99 %v", p99(bulk), p99(direct)) + Expect(p99(shared)).To(BeNumerically(">", p99(bulk)), + "shared p99 %v, bulk p99 %v", p99(shared), p99(bulk)) + }) + }) +}) diff --git a/core/services/worker/tunnel_size_race_test.go b/core/services/worker/tunnel_size_race_test.go new file mode 100644 index 000000000000..7fe838546686 --- /dev/null +++ b/core/services/worker/tunnel_size_race_test.go @@ -0,0 +1,5 @@ +//go:build race + +package worker + +const transferBytes = 128 << 20 diff --git a/core/services/worker/tunnel_size_test.go b/core/services/worker/tunnel_size_test.go new file mode 100644 index 000000000000..0f2a4a94c1b0 --- /dev/null +++ b/core/services/worker/tunnel_size_test.go @@ -0,0 +1,8 @@ +//go:build !race + +package worker + +// transferBytes is the size of the transfer in the specs that measure the delay +// of a small call. The race detector slows the copy so much that a build with it +// moves less (see tunnel_size_race_test.go). +const transferBytes = 256 << 20 diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go new file mode 100644 index 000000000000..3aa10e8c14e7 --- /dev/null +++ b/core/services/worker/tunnel_test.go @@ -0,0 +1,946 @@ +package worker + +import ( + "context" + "encoding/binary" + "errors" + "fmt" + "io" + "math/rand/v2" + "net" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// awaitErr runs fn on its own goroutine and reports its result on a channel. +// +// Every blocking read in this file goes through it. To assert that the worker +// refused a stream promptly by arming a read deadline and expecting an error +// proves nothing, because a stream that the worker never answers satisfies the +// deadline as well as one that it refused. A read with no deadline on another +// goroutine, and an Eventually on the channel, turns that round: a stream that +// is parked delivers nothing, and the spec fails. +func awaitErr(fn func() error) <-chan error { + ch := make(chan error, 1) + go func() { ch <- fn() }() + return ch +} + +// tunnelDial is what the fake frontend saw on one incoming dial. +type tunnelDial struct { + token string + nodeID string + lane tunnel.Lane +} + +// fakeFrontend is the far side of the tunnel. It speaks the real websocket +// upgrade and the real session handshake of the tunnel package, so that these +// specs exercise the wire and not a mock of it. It is not the connect endpoint: +// that one needs a database, and the client is what is under test here. +type fakeFrontend struct { + srv *httptest.Server + inference chan *tunnel.Session + bulk chan *tunnel.Session + dials chan tunnelDial + + // closeAtOnce makes every accepted session end at once, which is what a + // replica that goes down in a rolling restart looks like to the worker. + closeAtOnce bool + // bulkStatus, when it is not zero, is the status that a bulk dial gets + // instead of an upgrade. + bulkStatus atomic.Int32 +} + +func newFakeFrontend(closeAtOnce bool) *fakeFrontend { + f := &fakeFrontend{ + inference: make(chan *tunnel.Session, 64), + bulk: make(chan *tunnel.Session, 64), + dials: make(chan tunnelDial, 256), + closeAtOnce: closeAtOnce, + } + upgrader := tunnel.NewUpgrader() + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != tunnel.ConnectPath { + w.WriteHeader(http.StatusNotFound) + return + } + lane, err := tunnel.ParseLane(r.URL.Query().Get("lane")) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + select { + case f.dials <- tunnelDial{ + token: strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer "), + nodeID: r.URL.Query().Get("id"), + lane: lane, + }: + default: + } + if status := int(f.bulkStatus.Load()); lane == tunnel.LaneBulk && status != 0 { + w.WriteHeader(status) + return + } + + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := tunnel.ServerSession(ws, lane) + if err != nil { + _ = ws.Close() + return + } + if f.closeAtOnce { + _ = sess.Close() + return + } + sessions := f.inference + if lane == tunnel.LaneBulk { + sessions = f.bulk + } + select { + case sessions <- sess: + default: + _ = sess.Close() + } + })) + return f +} + +func (f *fakeFrontend) close() { + for _, sessions := range []chan *tunnel.Session{f.inference, f.bulk} { + drain: + for { + select { + case sess := <-sessions: + _ = sess.Close() + default: + break drain + } + } + } + f.srv.Close() +} + +// echoListenerOn starts a TCP server that writes back what it reads. +func echoListenerOn(addr string) net.Listener { + ln, err := net.Listen("tcp", addr) + Expect(err).ToNot(HaveOccurred()) + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go func() { + defer func() { _ = conn.Close() }() + _, _ = io.Copy(conn, conn) + }() + } + }() + return ln +} + +func echoListener() net.Listener { return echoListenerOn("127.0.0.1:0") } + +// portOf returns the port that a listener is bound to. +func portOf(ln net.Listener) int { + _, portStr, err := net.SplitHostPort(ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + port, err := strconv.Atoi(portStr) + Expect(err).ToNot(HaveOccurred()) + return port +} + +// dialLocalTCP is the simplest LocalService: connect to what the frontend named. +func dialLocalTCP(ctx context.Context, target string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", target) +} + +var _ = Describe("Worker tunnel client", func() { + var ( + ctx context.Context + cancel context.CancelFunc + frontend *fakeFrontend + tun *Tunnel + ) + + BeforeEach(func() { + ctx, cancel = context.WithCancel(context.Background()) + }) + + AfterEach(func() { + if tun != nil { + Expect(tun.Close()).To(Succeed()) + tun = nil + } + cancel() + if frontend != nil { + frontend.close() + frontend = nil + } + }) + + start := func(mutate func(*TunnelConfig)) { + GinkgoHelper() + cfg := TunnelConfig{ + FrontendURL: frontend.srv.URL, + NodeID: "node-1", + Token: func() string { return "tunnel-secret" }, + Services: map[string]LocalService{}, + } + if mutate != nil { + mutate(&cfg) + } + var err error + tun, err = StartTunnel(ctx, cfg) + Expect(err).ToNot(HaveOccurred()) + } + + session := func() *tunnel.Session { + GinkgoHelper() + var sess *tunnel.Session + Eventually(frontend.inference, "10s").Should(Receive(&sess)) + return sess + } + + // open opens a stream, names the service, and waits for the reply. + open := func(sess *tunnel.Session, tag, target string) (net.Conn, error) { + GinkgoHelper() + stream, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(stream, tag, target)).To(Succeed()) + var got error + Eventually(awaitErr(func() error { return tunnel.ReadStreamReply(stream) }), "10s").Should(Receive(&got)) + return stream, got + } + + echoes := func(stream net.Conn, msg string) { + GinkgoHelper() + _, err := stream.Write([]byte(msg)) + Expect(err).ToNot(HaveOccurred()) + buf := make([]byte, len(msg)) + Eventually(awaitErr(func() error { _, err := io.ReadFull(stream, buf); return err }), "10s").Should(Receive(BeNil())) + Expect(string(buf)).To(Equal(msg)) + } + + Describe("carrying a tagged stream to a local service", func() { + It("routes a stream tagged for gRPC to the local address it names", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { c.Services[tunnel.StreamTagGRPC] = dialLocalTCP }) + + stream, err := open(session(), tunnel.StreamTagGRPC, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + echoes(stream, "ping") + }) + + It("routes through the own table of the worker and ignores the host the frontend names", func() { + // The other specs install dialLocalTCP, which dials whatever it gets. + // This one installs tunnelServices, the table that Run installs, so + // the wire is tried at least once against the real routing rules. + backend := echoListener() + DeferCleanup(func() { _ = backend.Close() }) + port := portOf(backend) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services = tunnelServices(&Config{ServeAddr: fmt.Sprintf("0.0.0.0:%d", port), GRPCMaxPort: port}, "0.0.0.0:1") + }) + + stream, err := open(session(), tunnel.StreamTagGRPC, fmt.Sprintf("attacker.invalid:%d", port)) + Expect(err).ToNot(HaveOccurred()) + echoes(stream, "loopback") + }) + + It("refuses a port outside its range as a bad request, over the wire", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services = tunnelServices(&Config{ServeAddr: "0.0.0.0:50051", GRPCMaxPort: 50051}, "0.0.0.0:50050") + }) + + _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:22") + Expect(err).To(MatchError(tunnel.ErrStreamRequestInvalid)) + Expect(err).ToNot(MatchError(tunnel.ErrStreamTargetUnavailable)) + Expect(err).ToNot(MatchError(tunnel.ErrStreamTagUnknown)) + }) + + It("routes the http tag to the own server of the worker and ignores the target", func() { + own := echoListener() + DeferCleanup(func() { _ = own.Close() }) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services = tunnelServices(&Config{}, fmt.Sprintf("0.0.0.0:%d", portOf(own))) + }) + + stream, err := open(session(), tunnel.StreamTagHTTP, "attacker.invalid:1") + Expect(err).ToNot(HaveOccurred()) + echoes(stream, "own-server") + }) + }) + + Describe("refusing a stream it cannot serve", func() { + It("refuses an unknown tag promptly, and the stream ends and does not park", func() { + frontend = newFakeFrontend(false) + start(nil) + + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(stream, "no-such-tag", "")).To(Succeed()) + + // Two facts in order on one goroutine: the worker said why, and the + // stream then ended. A worker that says why and leaves the stream + // open never sends on the channel. + type outcome struct{ reply, end error } + done := make(chan outcome, 1) + go func() { + var got outcome + got.reply = tunnel.ReadStreamReply(stream) + _, got.end = stream.Read(make([]byte, 1)) + done <- got + }() + var got outcome + Eventually(done, "10s").Should(Receive(&got)) + Expect(got.reply).To(MatchError(tunnel.ErrStreamTagUnknown)) + Expect(got.end).To(MatchError(io.EOF)) + }) + + It("reports a local service it could not reach as unavailable and not as an unknown tag", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { + return nil, errors.New("connection refused") + } + }) + _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:1") + Expect(err).To(MatchError(tunnel.ErrStreamTargetUnavailable)) + Expect(err).ToNot(MatchError(tunnel.ErrStreamTagUnknown)) + }) + + It("ends a stream whose request never arrives and does not hold it open", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = dialLocalTCP + c.headerTimeout = 50 * time.Millisecond + }) + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + ended := awaitErr(func() error { _, err := io.Copy(io.Discard, stream); return err }) + Eventually(ended, "10s").Should(Receive(BeNil())) + }) + + It("says that it learned nothing when the request frame did not arrive in time", func() { + // A frame that arrives late is not a frame that is malformed. On the + // relay path a round trip of the link between replicas runs inside + // this window, so a refusal that counted as evidence would evict a + // model across the fleet for nothing but a congested link. + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = dialLocalTCP + c.headerTimeout = 50 * time.Millisecond + }) + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + var got error + Eventually(awaitErr(func() error { return tunnel.ReadStreamReply(stream) }), "10s").Should(Receive(&got)) + + Expect(got).To(MatchError(tunnel.ErrStreamNotServed)) + Expect(got).ToNot(MatchError(tunnel.ErrStreamRequestInvalid)) + Expect(tunnel.IsWorkerAnswer(got)).To(BeFalse()) + }) + + It("still calls a malformed request frame malformed, which is evidence", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = dialLocalTCP + c.headerTimeout = time.Minute + }) + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(binary.Write(stream, binary.BigEndian, uint16(60000))).To(Succeed()) + var got error + Eventually(awaitErr(func() error { return tunnel.ReadStreamReply(stream) }), "10s").Should(Receive(&got)) + + Expect(got).To(MatchError(tunnel.ErrStreamRequestInvalid)) + Expect(tunnel.IsWorkerAnswer(got)).To(BeTrue()) + }) + + It("says that it learned nothing when the local dial ended because the session went away", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { + return nil, fmt.Errorf("dialing the backend: %w", context.Canceled) + } + }) + _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:41000") + Expect(err).To(MatchError(tunnel.ErrStreamNotServed)) + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("does not count a timeout of the local dial as the answer of the target", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { + return nil, &net.OpError{Op: "dial", Err: timeoutError{}} + } + }) + _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:41000") + Expect(err).To(MatchError(tunnel.ErrStreamNotServed)) + }) + + DescribeTable("keeps a classification that the local service already made", + func(classified error, wantEvidence bool) { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { + return nil, fmt.Errorf("%w: from the service", classified) + } + }) + _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:1") + Expect(err).To(MatchError(classified)) + Expect(tunnel.IsWorkerAnswer(err)).To(Equal(wantEvidence)) + }, + Entry("I learned nothing", tunnel.ErrStreamNotServed, false), + Entry("I do not serve that tag", tunnel.ErrStreamTagUnknown, true), + Entry("that request was malformed", tunnel.ErrStreamRequestInvalid, true), + Entry("I could not reach the target", tunnel.ErrStreamTargetUnavailable, true), + ) + }) + + Describe("surviving a bad stream", func() { + It("keeps serving the session after a stream it could not read", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { c.Services[tunnel.StreamTagGRPC] = dialLocalTCP }) + sess := session() + + // A frame that declares much more than it sends, and then hangs up. + bad, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + var hdr [2]byte + binary.BigEndian.PutUint16(hdr[:], 900) + _, err = bad.Write(append(hdr[:], []byte("gr")...)) + Expect(err).ToNot(HaveOccurred()) + Expect(bad.(interface{ CloseWrite() error }).CloseWrite()).To(Succeed()) + Eventually(awaitErr(func() error { _, err := io.Copy(io.Discard, bad); return err }), "10s").Should(Receive(BeNil())) + + good, err := open(sess, tunnel.StreamTagGRPC, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + echoes(good, "still here") + }) + + It("serves streams at the same time, so a stream that is live does not block the next", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { c.Services[tunnel.StreamTagGRPC] = dialLocalTCP }) + sess := session() + + _, err := open(sess, tunnel.StreamTagGRPC, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + _, err = open(sess, tunnel.StreamTagGRPC, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + }) + + It("keeps the session after a local service panics", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services["explodes"] = func(context.Context, string) (net.Conn, error) { panic("a local service blew up") } + c.Services[tunnel.StreamTagGRPC] = dialLocalTCP + }) + sess := session() + + boom, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(boom, "explodes", "")).To(Succeed()) + Eventually(awaitErr(func() error { _, err := io.Copy(io.Discard, boom); return err }), "10s").Should(Receive(BeNil())) + + _, err = open(sess, tunnel.StreamTagGRPC, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + }) + }) + + Describe("connecting again", func() { + recordDelays := func(delays chan time.Duration) func(*TunnelConfig) { + return func(c *TunnelConfig) { + c.noBulk = true + c.sleep = func(ctx context.Context, d time.Duration) error { + select { + case delays <- d: + default: + } + return ctx.Err() + } + } + } + + It("waits longer after each attempt, bounded and never in a tight loop", func() { + frontend = newFakeFrontend(true) // every session ends at once + delays := make(chan time.Duration, 64) + start(recordDelays(delays)) + + observed := make([]time.Duration, 0, 10) + for i := range 10 { + var d time.Duration + Eventually(delays, "20s").Should(Receive(&d), "attempt %d", i+1) + observed = append(observed, d) + } + for i, d := range observed { + Expect(d).To(BeNumerically(">", 0), "delay %d", i+1) + Expect(d).To(BeNumerically("<=", tunnelBackoffMax), "delay %d", i+1) + } + // The jitter keeps half of the delay, so the fourth wait is at least + // four times the base, whatever the dice say. + Expect(observed[3]).To(BeNumerically(">=", 4*tunnelBackoffBase)) + }) + + It("does not return to the shortest wait after a session that ended at once", func() { + frontend = newFakeFrontend(true) + delays := make(chan time.Duration, 64) + start(recordDelays(delays)) + var last time.Duration + for range 6 { + Eventually(delays, "20s").Should(Receive(&last)) + } + Expect(last).To(BeNumerically(">", tunnelBackoffBase)) + }) + + It("returns to the shortest wait after a session that lasted", func() { + frontend = newFakeFrontend(true) + // A clock that jumps one minute at every reading. The loop reads it + // when a session starts and when it ends, so every session seems to + // have lasted a minute. + var ticks atomic.Int64 + base := time.Now() + delays := make(chan time.Duration, 64) + start(func(c *TunnelConfig) { + recordDelays(delays)(c) + c.now = func() time.Time { return base.Add(time.Duration(ticks.Add(1)) * time.Minute) } + }) + for i := range 6 { + var d time.Duration + Eventually(delays, "20s").Should(Receive(&d), "attempt %d", i+1) + Expect(d).To(BeNumerically("<=", tunnelBackoffBase), "delay %d", i+1) + } + }) + + It("presents the credential that is current at the time of the dial", func() { + frontend = newFakeFrontend(true) + var issued atomic.Int64 + start(func(c *TunnelConfig) { + c.Token = func() string { return fmt.Sprintf("token-%d", issued.Add(1)) } + c.sleep = func(ctx context.Context, _ time.Duration) error { return ctx.Err() } + }) + var first, second tunnelDial + Eventually(frontend.dials, "20s").Should(Receive(&first)) + Eventually(frontend.dials, "20s").Should(Receive(&second)) + Expect(first.token).To(Equal("token-1")) + Expect(second.token).To(Equal("token-2")) + Expect(first.nodeID).To(Equal("node-1")) + }) + + It("does not dial with no credential, and says so", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Token = func() string { return "" } + c.sleep = func(ctx context.Context, _ time.Duration) error { <-ctx.Done(); return ctx.Err() } + }) + Consistently(frontend.dials, "500ms", "50ms").ShouldNot(Receive()) + Expect(tun.Connected()).To(BeFalse()) + }) + }) + + Describe("the two lanes", func() { + It("dials the bulk lane after the inference lane, with the lane in the query", func() { + frontend = newFakeFrontend(false) + start(nil) + session() + + var first, second tunnelDial + Eventually(frontend.dials, "10s").Should(Receive(&first)) + Eventually(frontend.dials, "10s").Should(Receive(&second)) + Expect(first.lane).To(Equal(tunnel.LaneInference)) + Expect(second.lane).To(Equal(tunnel.LaneBulk)) + Expect(second.token).To(Equal("tunnel-secret")) + Expect(second.nodeID).To(Equal("node-1")) + Eventually(tun.BulkConnected, "10s").Should(BeTrue()) + }) + + It("serves the streams of the bulk session as it serves those of the inference session", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { c.Services[tunnel.StreamTagHTTP] = dialLocalTCP }) + session() + + var bulk *tunnel.Session + Eventually(frontend.bulk, "10s").Should(Receive(&bulk)) + stream, err := open(bulk, tunnel.StreamTagHTTP, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + echoes(stream, "file bytes") + }) + + It("keeps the inference lane when the frontend refuses the bulk lane", func() { + for _, status := range []int{http.StatusNotFound, http.StatusForbidden, http.StatusServiceUnavailable, http.StatusConflict} { + frontend = newFakeFrontend(false) + frontend.bulkStatus.Store(int32(status)) + start(nil) + session() + Eventually(frontend.dials, "10s").Should(Receive()) // the inference dial + Eventually(frontend.dials, "10s").Should(Receive(), "a bulk dial that was refused with %d is tried again", status) + Expect(tun.Connected()).To(BeTrue(), "status %d", status) + Expect(tun.BulkConnected()).To(BeFalse(), "status %d", status) + + Expect(tun.Close()).To(Succeed()) + tun = nil + frontend.close() + frontend = nil + } + }) + + It("dials the bulk lane again soon after the answer 409, because another replica may hold the tunnel", func() { + frontend = newFakeFrontend(false) + frontend.bulkStatus.Store(http.StatusConflict) + delays := make(chan time.Duration, 64) + start(func(c *TunnelConfig) { + c.sleep = func(ctx context.Context, d time.Duration) error { + select { + case delays <- d: + default: + } + // Do not wait. The delay is what is measured. + return ctx.Err() + } + }) + session() + for range 5 { + var d time.Duration + Eventually(delays, "10s").Should(Receive(&d)) + Expect(d).To(BeNumerically("<=", tunnelWrongReplicaDelay*3/2), "a 409 must not grow the wait") + } + + // The right replica answers now. + frontend.bulkStatus.Store(0) + Eventually(tun.BulkConnected, "10s").Should(BeTrue()) + }) + + It("connects both lanes again after the inference session ends", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { c.sleep = func(ctx context.Context, _ time.Duration) error { return ctx.Err() } }) + inference := session() + Eventually(tun.BulkConnected, "10s").Should(BeTrue()) + + Expect(inference.Close()).To(Succeed()) + + Eventually(func() bool { return tun.Connected() && tun.BulkConnected() }, "20s").Should(BeTrue()) + Eventually(frontend.inference, "10s").Should(Receive()) + }) + + It("dials the bulk lane again when only the bulk session ends, and keeps the inference lane", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { c.sleep = func(ctx context.Context, _ time.Duration) error { return ctx.Err() } }) + inference := session() + var bulk *tunnel.Session + Eventually(frontend.bulk, "10s").Should(Receive(&bulk)) + Eventually(tun.BulkConnected, "10s").Should(BeTrue()) + + Expect(bulk.Close()).To(Succeed()) + + Eventually(frontend.bulk, "10s").Should(Receive(), "the worker dials the bulk lane again") + Expect(inference.IsClosed()).To(BeFalse()) + Expect(tun.Connected()).To(BeTrue()) + }) + + It("closes the bulk session with the tunnel", func() { + frontend = newFakeFrontend(false) + start(nil) + session() + var bulk *tunnel.Session + Eventually(frontend.bulk, "10s").Should(Receive(&bulk)) + + Expect(tun.Close()).To(Succeed()) + tun = nil + Eventually(bulk.CloseChan(), "10s").Should(BeClosed()) + }) + }) + + Describe("reporting whether it holds a session", func() { + It("reports connected once the frontend has accepted its dial", func() { + frontend = newFakeFrontend(false) + start(nil) + session() + Eventually(tun.Connected, "10s").Should(BeTrue()) + }) + + It("reports disconnected once the session is gone", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + // The next dial waits, so that the spec sees the gap. + c.sleep = func(ctx context.Context, _ time.Duration) error { <-ctx.Done(); return ctx.Err() } + }) + sess := session() + Eventually(tun.Connected, "10s").Should(BeTrue()) + Expect(sess.Close()).To(Succeed()) + Eventually(tun.Connected, "10s").Should(BeFalse()) + }) + + It("reports disconnected before the first dial has landed", func() { + frontend = newFakeFrontend(false) + frontend.srv.Close() + start(func(c *TunnelConfig) { + c.sleep = func(ctx context.Context, _ time.Duration) error { <-ctx.Done(); return ctx.Err() } + }) + Consistently(tun.Connected, "500ms", "50ms").Should(BeFalse()) + }) + + It("reports disconnected on a tunnel that is nil", func() { + var absent *Tunnel + Expect(absent.Connected()).To(BeFalse()) + Expect(absent.BulkConnected()).To(BeFalse()) + }) + + It("is what the readiness gate of the worker looks at", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.sleep = func(ctx context.Context, _ time.Duration) error { <-ctx.Done(); return ctx.Err() } + }) + readiness := &nodes.WorkerReadiness{} + readiness.Set(nodes.TunnelReadiness(tun)) + sess := session() + Eventually(readiness.Check, "10s").Should(Succeed()) + // The gate fails open when no probe is installed, so a gate that + // goes back to not ready proves that the probe is installed. + Expect(sess.Close()).To(Succeed()) + Eventually(readiness.Check, "10s").Should(MatchError(nodes.ErrTunnelDisconnected)) + }) + }) + + Describe("starting", func() { + It("refuses a configuration that cannot work, and not a frontend that cannot be reached", func() { + _, err := StartTunnel(ctx, TunnelConfig{FrontendURL: "http://frontend:8080", Token: func() string { return "t" }}) + Expect(err).To(MatchError(ContainSubstring("no node id"))) + _, err = StartTunnel(ctx, TunnelConfig{FrontendURL: "http://frontend:8080", NodeID: "n"}) + Expect(err).To(MatchError(ContainSubstring("no credential source"))) + _, err = StartTunnel(ctx, TunnelConfig{FrontendURL: "ftp://frontend", NodeID: "n", Token: func() string { return "t" }}) + Expect(err).To(HaveOccurred()) + + // A frontend that is down is not a reason to refuse the start. + down := httptest.NewServer(http.NotFoundHandler()) + url := down.URL + down.Close() + t, err := StartTunnel(ctx, TunnelConfig{FrontendURL: url, NodeID: "n", Token: func() string { return "t" }, + sleep: func(ctx context.Context, _ time.Duration) error { <-ctx.Done(); return ctx.Err() }}) + Expect(err).ToNot(HaveOccurred()) + Expect(t.Close()).To(Succeed()) + }) + + It("can be closed twice", func() { + frontend = newFakeFrontend(false) + start(nil) + Expect(tun.Close()).To(Succeed()) + Expect(tun.Close()).To(Succeed()) + tun = nil + }) + }) +}) + +// timeoutError is a net.Error that says that it is a timeout. +type timeoutError struct{} + +func (timeoutError) Error() string { return "i/o timeout" } +func (timeoutError) Timeout() bool { return true } +func (timeoutError) Temporary() bool { return true } + +// listenOnSecondLoopback binds 127.0.0.2 on a port where nothing listens on +// 127.0.0.1. +// +// The choice of the port is part of the assertion that this serves. Only +// 127.0.0.2 listens, so a service that follows the host the frontend names would +// connect, and one that dials loopback cannot. A port from :0 is in the +// ephemeral range, where an unrelated socket on 127.0.0.1 can hold the same +// number, and then the dial to 127.0.0.1 succeeds and the spec reports a hole +// that is not there. A port below the ephemeral range (32768 by default on +// Linux) is never given out for outbound connections. 127.0.0.1 is tried and +// released and not held, because holding it would make the dial that must fail +// succeed. +func listenOnSecondLoopback() (net.Listener, int) { + GinkgoHelper() + const ( + floor = 20000 + ceiling = 31000 + attempts = 200 + ) + for i := range attempts { + // #nosec G404 -- a port for a test. + port := floor + rand.IntN(ceiling-floor) + free, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + continue + } + if err := free.Close(); err != nil { + continue + } + victim, err := net.Listen("tcp", fmt.Sprintf("127.0.0.2:%d", port)) + if err != nil { + if i == 0 && strings.Contains(err.Error(), "assign requested address") { + Skip("this host cannot bind a second loopback address: " + err.Error()) + } + continue + } + return victim, port + } + Fail(fmt.Sprintf("no port in [%d, %d) was free on 127.0.0.1 and bindable on 127.0.0.2 after %d attempts", floor, ceiling, attempts)) + return nil, 0 +} + +// The routing table is the security boundary of the tunnel. Every spec above +// that uses dialLocalTCP installs the dialler that loopbackService exists to +// prevent, so these specs test the real functions. +var _ = Describe("Worker tunnel local services", func() { + var ctx context.Context + + BeforeEach(func() { ctx = context.Background() }) + + Describe("loopbackService", func() { + It("reaches a loopback listener whose port is in range", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + port := portOf(ln) + conn, err := loopbackService(port, port)(ctx, fmt.Sprintf("127.0.0.1:%d", port)) + Expect(err).ToNot(HaveOccurred()) + Expect(conn.Close()).To(Succeed()) + }) + + It("ignores the host the frontend names and dials loopback", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + port := portOf(ln) + conn, err := loopbackService(port, port)(ctx, fmt.Sprintf("attacker.invalid:%d", port)) + Expect(err).ToNot(HaveOccurred()) + Expect(conn.Close()).To(Succeed()) + }) + + It("does not reach a listener on another local address that the frontend names", func() { + victim, port := listenOnSecondLoopback() + DeferCleanup(func() { _ = victim.Close() }) + _, err := loopbackService(port, port)(ctx, fmt.Sprintf("127.0.0.2:%d", port)) + Expect(err).To(HaveOccurred()) + }) + + DescribeTable("refuses a target it will not route", + func(target string, minPort, maxPort int) { + _, err := loopbackService(minPort, maxPort)(ctx, target) + Expect(err).To(MatchError(tunnel.ErrStreamRequestInvalid)) + }, + Entry("a port below the range", "127.0.0.1:50050", 50051, 50060), + Entry("a port above the range", "127.0.0.1:50061", 50051, 50060), + Entry("a port that is not a number", "127.0.0.1:http", 50051, 50060), + Entry("no port", "127.0.0.1", 50051, 50060), + Entry("an empty target", "", 50051, 50060), + ) + + It("reports a backend that does not listen as unavailable", func() { + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + port := portOf(ln) + Expect(ln.Close()).To(Succeed()) + + _, err = loopbackService(port, port)(ctx, fmt.Sprintf("127.0.0.1:%d", port)) + Expect(err).To(HaveOccurred()) + Expect(classifyServiceFailure(err)).To(MatchError(tunnel.ErrStreamTargetUnavailable)) + }) + }) + + Describe("fixedService", func() { + It("reaches its own address whatever the frontend names", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + conn, err := fixedService(ln.Addr().String())(ctx, "attacker.invalid:1") + Expect(err).ToNot(HaveOccurred()) + Expect(conn.Close()).To(Succeed()) + }) + }) + + DescribeTable("loopbackAddr turns a bind address into one that can be dialled", + func(bind, want string) { Expect(loopbackAddr(bind)).To(Equal(want)) }, + Entry("IPv4 wildcard", "0.0.0.0:8080", "127.0.0.1:8080"), + Entry("IPv6 wildcard", "[::]:8080", "127.0.0.1:8080"), + Entry("no host", ":8080", "127.0.0.1:8080"), + Entry("an explicit host is left alone", "10.0.0.9:8080", "10.0.0.9:8080"), + Entry("an explicit loopback is left alone", "127.0.0.1:8080", "127.0.0.1:8080"), + Entry("something that is not host:port passes through", "not-an-address", "not-an-address"), + ) + + Describe("tunnelServices", func() { + It("serves exactly the two tags that the frontend may name", func() { + table := tunnelServices(&Config{}, "0.0.0.0:50050") + Expect(table).To(HaveLen(2)) + Expect(table).To(HaveKey(tunnel.StreamTagGRPC)) + Expect(table).To(HaveKey(tunnel.StreamTagHTTP)) + }) + + It("limits the gRPC service to the port range of this worker", func() { + table := tunnelServices(&Config{ServeAddr: "0.0.0.0:50051", GRPCMaxPort: 50055}, "0.0.0.0:50050") + _, err := table[tunnel.StreamTagGRPC](ctx, "127.0.0.1:50056") + Expect(err).To(MatchError(tunnel.ErrStreamRequestInvalid)) + _, err = table[tunnel.StreamTagGRPC](ctx, "127.0.0.1:50050") + Expect(err).To(MatchError(tunnel.ErrStreamRequestInvalid)) + }) + }) + + DescribeTable("tunnelEndpoint builds the URL that the worker dials", + func(frontend, nodeID, want string) { + got, err := tunnelEndpoint(frontend, nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal(want)) + }, + Entry("http becomes ws", "http://frontend:8080", "n1", "ws://frontend:8080/api/cluster/connect?id=n1"), + Entry("https becomes wss", "https://frontend", "n1", "wss://frontend/api/cluster/connect?id=n1"), + Entry("ws stays", "ws://frontend:8080", "n1", "ws://frontend:8080/api/cluster/connect?id=n1"), + Entry("wss stays", "wss://frontend", "n1", "wss://frontend/api/cluster/connect?id=n1"), + Entry("a path prefix stays", "https://host/localai", "n1", "wss://host/localai/api/cluster/connect?id=n1"), + Entry("a slash at the end is not doubled", "https://host/localai/", "n1", "wss://host/localai/api/cluster/connect?id=n1"), + Entry("the node id is escaped", "http://h", "a b&c", "ws://h/api/cluster/connect?id=a+b%26c"), + ) + + DescribeTable("tunnelEndpoint refuses a URL it cannot dial", + func(frontend string) { + _, err := tunnelEndpoint(frontend, "n1") + Expect(err).To(HaveOccurred()) + }, + Entry("empty", ""), + Entry("a scheme that is not HTTP", "ftp://frontend"), + Entry("a host with no scheme", "frontend:8080/x"), + Entry("no host", "http://"), + ) + + DescribeTable("tunnelBackoffDelay stays between half of the step and the step, and below the limit", + func(attempt int) { + for range 200 { + d := tunnelBackoffDelay(attempt) + Expect(d).To(BeNumerically(">", 0)) + Expect(d).To(BeNumerically("<=", tunnelBackoffMax)) + } + }, + Entry("zero", 0), + Entry("one", 1), + Entry("six", 6), + Entry("forty", 40), + Entry("a number that would overflow the shift", 1000), + ) +}) From 06f10410827f815253796fba18bcc7d8f282f7c6 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:50:47 +0000 Subject: [PATCH 028/162] test(tunnel): require the shared constructors for every websocket A spec reads the files of the connect endpoint and of the worker client and fails if one of them builds a gorilla upgrader or dialer by hand. Such a websocket gets the 4 KiB buffers, and a transfer from the worker runs five times slower than in the other direction. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/session_test.go | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/core/services/tunnel/session_test.go b/core/services/tunnel/session_test.go index 11cc3eb536a5..509de8e61728 100644 --- a/core/services/tunnel/session_test.go +++ b/core/services/tunnel/session_test.go @@ -3,6 +3,10 @@ package tunnel_test import ( "context" "io" + "os" + "path/filepath" + "slices" + "strings" "time" "github.com/mudler/LocalAI/core/services/tunnel" @@ -23,6 +27,27 @@ var _ = Describe("Buffers of the websocket", func() { Expect(d.HandshakeTimeout).To(Equal(3 * time.Second)) }) + It("is the only way the files of the tunnel make a websocket", func() { + // A literal websocket.Upgrader or websocket.Dialer gets the 4 KiB buffers + // of gorilla, and a transfer from the worker to the frontend runs five + // times slower than the other direction. The files that open a websocket + // of the tunnel must use NewUpgrader and NewDialer. + var files []string + for _, pattern := range []string{"../../http/endpoints/cluster/*.go", "../worker/tunnel*.go"} { + matches, err := filepath.Glob(pattern) + Expect(err).ToNot(HaveOccurred()) + files = append(files, matches...) + } + files = slices.DeleteFunc(files, func(f string) bool { return strings.HasSuffix(f, "_test.go") }) + Expect(files).ToNot(BeEmpty(), "the pattern found no file, so it checks nothing") + for _, f := range files { + src, err := os.ReadFile(f) + Expect(err).ToNot(HaveOccurred()) + Expect(string(src)).ToNot(ContainSubstring("websocket.Upgrader{"), f) + Expect(string(src)).ToNot(ContainSubstring("websocket.Dialer{"), f) + } + }) + It("keeps the default origin check so that a browser of another origin is refused", func() { Expect(tunnel.NewUpgrader().CheckOrigin).To(BeNil()) }) From 73934c1be7e8ece309bada38be42c2e1ac545274 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:50:47 +0000 Subject: [PATCH 029/162] feat(worker): attach to the carrier that the frontend names The worker registers first and learns the carrier from the answer. For the tunnel it starts the tunnel client and needs no NATS URL. For NATS, or for a frontend that names no carrier, it connects to NATS as before, and it stops with a message that says so if it has no URL. The flag is no longer required, because for a cluster on the tunnel nothing needs it. The heartbeat and the readiness probe look at the link of the carrier that the worker is attached to. A worker whose link is down does not send its heartbeat, so the registry does not show a node that looks healthy and cannot be reached. A worker that is reached only through its tunnel and was started with no address binds its backends and its file-transfer server to loopback, and tells the frontend that its address is not routable. A worker with an address, or any worker on NATS, binds all interfaces as before. The lifecycle verbs have no route on the tunnel yet, and the worker says so when it starts. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/addr_test.go | 32 +++- core/services/worker/carrier_test.go | 27 ++++ core/services/worker/config.go | 9 +- core/services/worker/nats_connect.go | 34 +++++ core/services/worker/registration.go | 37 ++++- core/services/worker/supervisor.go | 15 +- core/services/worker/worker.go | 210 ++++++++++++++++++--------- 7 files changed, 283 insertions(+), 81 deletions(-) create mode 100644 core/services/worker/carrier_test.go diff --git a/core/services/worker/addr_test.go b/core/services/worker/addr_test.go index 4f5b1ba67f4f..6876e64d8ff0 100644 --- a/core/services/worker/addr_test.go +++ b/core/services/worker/addr_test.go @@ -57,7 +57,7 @@ var _ = Describe("Worker address resolution", func() { DescribeTable("returns the correct address", func(httpAddr, addr, serve, want string) { cfg := &Config{HTTPAddr: httpAddr, Addr: addr, ServeAddr: serve} - Expect(cfg.resolveHTTPAddr()).To(Equal(want)) + Expect(cfg.resolveHTTPAddr("0.0.0.0")).To(Equal(want)) }, Entry("HTTPAddr takes priority", "0.0.0.0:8080", "", "", "0.0.0.0:8080"), Entry("derives from Addr port minus 1", "", "worker1:60000", "0.0.0.0:50051", "0.0.0.0:59999"), @@ -66,6 +66,36 @@ var _ = Describe("Worker address resolution", func() { ) }) + Describe("the bind host", func() { + DescribeTable("is loopback only for a worker on the tunnel with no routable address", + func(tunnel bool, addr, advertise, want string) { + cfg := &Config{Addr: addr, AdvertiseAddr: advertise} + Expect(cfg.bindHost(tunnel)).To(Equal(want)) + }, + Entry("tunnel, no address", true, "", "", "127.0.0.1"), + Entry("tunnel, --addr", true, "worker1:50051", "", "0.0.0.0"), + Entry("tunnel, --advertise-addr", true, "", "worker1:50051", "0.0.0.0"), + Entry("NATS, no address: binds as it always did", false, "", "", "0.0.0.0"), + Entry("NATS, --addr", false, "worker1:50051", "", "0.0.0.0"), + ) + + It("puts the host in front of the port of the file-transfer server", func() { + cfg := &Config{ServeAddr: "0.0.0.0:50051"} + Expect(cfg.resolveHTTPAddr("127.0.0.1")).To(Equal("127.0.0.1:50050")) + }) + + It("keeps an address that the operator set for the file-transfer server", func() { + cfg := &Config{HTTPAddr: "10.0.0.5:9000"} + Expect(cfg.resolveHTTPAddr("127.0.0.1")).To(Equal("10.0.0.5:9000")) + }) + + It("says in the registration if the address can be reached", func() { + Expect((&Config{}).registrationBody()["routable"]).To(BeFalse()) + Expect((&Config{Addr: "w:50051"}).registrationBody()["routable"]).To(BeTrue()) + Expect((&Config{AdvertiseAddr: "w:50051"}).registrationBody()["routable"]).To(BeTrue()) + }) + }) + Describe("advertiseHTTPAddr", func() { DescribeTable("returns the correct address", func(advertiseHTTP, advertise, addr, serve, want string) { diff --git a/core/services/worker/carrier_test.go b/core/services/worker/carrier_test.go new file mode 100644 index 000000000000..b9245328bcdf --- /dev/null +++ b/core/services/worker/carrier_test.go @@ -0,0 +1,27 @@ +package worker + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("The carrier that a worker attaches to", func() { + DescribeTable("follows the answer of the registration", + func(named, natsURL string, wantTunnel bool, wantErr string) { + onTunnel, err := carrierOf(named, &Config{NatsURL: natsURL}) + if wantErr != "" { + Expect(err).To(MatchError(ContainSubstring(wantErr))) + return + } + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(Equal(wantTunnel)) + }, + Entry("the tunnel, with no NATS URL", "tunnel", "", true, ""), + Entry("the tunnel, with a NATS URL that is not used", "tunnel", "nats://n:4222", true, ""), + Entry("NATS, with a URL", "nats", "nats://n:4222", false, ""), + Entry("no name (a frontend before carriers), with a URL", "", "nats://n:4222", false, ""), + Entry("NATS, with no URL", "nats", "", false, "no NATS URL"), + Entry("no name, with no URL", "", "", false, "does not name a carrier"), + Entry("a carrier from the future", "carrier-pigeon", "nats://n:4222", false, "does not know"), + ) +}) diff --git a/core/services/worker/config.go b/core/services/worker/config.go index 8a63ad8e0567..74f19c0e906e 100644 --- a/core/services/worker/config.go +++ b/core/services/worker/config.go @@ -9,7 +9,10 @@ package worker // Workers are backend-agnostic — they wait for backend.install NATS events // from the SmartRouter to install and start the required backend. // -// NATS is required. The worker acts as a process supervisor: +// The carrier is not a setting of the worker. The frontend names it in the answer +// to the registration. With NATS the worker needs the URL of the server. With the +// tunnel it needs nothing but the URL of the frontend. The worker acts as a +// process supervisor: // - Receives backend.install → installs backend from gallery, starts gRPC process, replies success // - Receives backend.stop → stops the gRPC process // - Receives stop → full shutdown (deregister + exit) @@ -75,8 +78,8 @@ type Config struct { // enforces it against the raw VRAM this worker reports. Empty = no cap. VRAMBudget string `env:"LOCALAI_VRAM_BUDGET" help:"Cap VRAM used for model allocation on this worker node, as a percentage (e.g. 80%) or absolute amount (e.g. 12GB)." group:"registration"` - // NATS (required) - NatsURL string `env:"LOCALAI_NATS_URL" required:"" help:"NATS server URL" group:"distributed"` + // NATS. The URL is needed only when the cluster runs on NATS. + NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed when the cluster runs on NATS; a cluster on the tunnel does not use it" group:"distributed"` NatsJWT string `env:"LOCALAI_NATS_JWT" help:"NATS user JWT override (normally from registration nats_jwt)" group:"distributed"` NatsUserSeed string `env:"LOCALAI_NATS_USER_SEED" help:"NATS user signing seed override (normally from registration nats_user_seed)" group:"distributed"` NatsRequireAuth bool `env:"LOCALAI_NATS_REQUIRE_AUTH" default:"false" help:"Require NATS JWT+seed from registration or env" group:"distributed"` diff --git a/core/services/worker/nats_connect.go b/core/services/worker/nats_connect.go index 25485701d2ec..a3db89de65af 100644 --- a/core/services/worker/nats_connect.go +++ b/core/services/worker/nats_connect.go @@ -1,9 +1,12 @@ package worker import ( + "context" "fmt" + "github.com/mudler/LocalAI/core/cli/workerregistry" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/xlog" ) // connectNATS opens a NATS client using JWT+seed from env or registration (env wins). @@ -31,3 +34,34 @@ func connectNATS(url, envJWT, envSeed, registerJWT, registerSeed string, require } return messaging.New(url, opts...) } + +// connectNatsFor connects a worker to NATS in the way it always did. +// +// Static credentials from the environment cannot be minted again, so they are +// used as they are. Otherwise the connection takes its credentials from the +// credential manager on every connect, and a goroutine refreshes them before the +// minted JWT expires. If the refresh fails for good, the worker stops, so that it +// restarts and acquires credentials again and does not drift towards a JWT that +// it can no longer renew. +func connectNatsFor(ctx context.Context, stop context.CancelFunc, cfg *Config, credMgr *workerregistry.CredentialManager, static bool, tls messaging.TLSFiles) (*messaging.Client, error) { + if static { + return connectNATS(cfg.NatsURL, cfg.NatsJWT, cfg.NatsUserSeed, "", "", cfg.NatsAuthRequired(), tls) + } + var opts []messaging.Option + if credMgr.HasCredentials() { + opts = append(opts, messaging.WithUserJWTProvider(credMgr.Provider())) + } + if tls.Enabled() { + opts = append(opts, messaging.WithTLS(tls)) + } + client, err := messaging.New(cfg.NatsURL, opts...) + if err == nil && credMgr.HasCredentials() { + go func() { + if err := credMgr.RefreshLoop(ctx); err != nil { + xlog.Error("NATS credential refresh permanently failed; shutting down worker", "error", err) + stop() + } + }() + } + return client, err +} diff --git a/core/services/worker/registration.go b/core/services/worker/registration.go index 23a439c724a1..a94cb03a0fa0 100644 --- a/core/services/worker/registration.go +++ b/core/services/worker/registration.go @@ -98,14 +98,38 @@ func (cfg *Config) advertiseAddr() string { return fmt.Sprintf("%s:%d", cmp.Or(hostname, "localhost"), cfg.effectiveBasePort()) } +// hasRoutableAddress reports whether the operator started this worker with an +// address that the frontends can reach (--addr or --advertise-addr). Such a +// worker can be reached by a frontend that dials, and it registers its address. +// A worker without one is reached only through its tunnel. +func (cfg *Config) hasRoutableAddress() bool { + return cfg.Addr != "" || cfg.AdvertiseAddr != "" +} + +// bindHost returns the host that this worker binds its backend processes and its +// file-transfer server to. +// +// A worker that is reached through its tunnel and has no routable address binds +// loopback. Nothing dials it, and an open port that nothing uses is surface. A +// worker with a routable address binds all interfaces, as a worker on NATS +// always did, because a frontend can dial it, and because it can follow a +// change of the carrier to NATS. A worker on NATS binds all interfaces whatever +// its address, so that a deployment on NATS behaves as before. +func (cfg *Config) bindHost(tunnelCarrier bool) string { + if tunnelCarrier && !cfg.hasRoutableAddress() { + return loopbackHost + } + return "0.0.0.0" +} + // resolveHTTPAddr returns the address to bind the HTTP file transfer server to. // Uses basePort-1 so it doesn't conflict with dynamically allocated gRPC ports // which grow upward from basePort. -func (cfg *Config) resolveHTTPAddr() string { +func (cfg *Config) resolveHTTPAddr(bindHost string) string { if cfg.HTTPAddr != "" { return cfg.HTTPAddr } - return fmt.Sprintf("0.0.0.0:%d", cfg.effectiveBasePort()-1) + return fmt.Sprintf("%s:%d", bindHost, cfg.effectiveBasePort()-1) } // advertiseHTTPAddr returns the HTTP address the frontend should use to reach @@ -165,9 +189,12 @@ func (cfg *Config) registrationBody() map[string]any { maxReplicas = 1 } body := map[string]any{ - "name": nodeName, - "address": cfg.advertiseAddr(), - "http_address": cfg.advertiseHTTPAddr(), + "name": nodeName, + "address": cfg.advertiseAddr(), + "http_address": cfg.advertiseHTTPAddr(), + // Says if the two addresses above can be reached. A frontend on the tunnel + // drops them when this is false. An older frontend ignores the key. + "routable": cfg.hasRoutableAddress(), "total_vram": totalVRAM, "available_vram": totalVRAM, // initially all VRAM is available "gpu_vendor": gpuVendor, diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 09c918d4b25c..3d7c86fcc07a 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -117,6 +117,11 @@ func (s *backendSupervisor) backendIdentity(name string) map[string]struct{} { // backendSupervisor manages multiple backend gRPC processes on different ports. // Each backend type (e.g., llama-cpp, bert-embeddings) gets its own process and port. type backendSupervisor struct { + // bindHost is the host that the backend processes listen on. It is empty for + // all interfaces. A worker that is reached only through its tunnel sets + // loopback. + bindHost string + cfg *Config ml *model.ModelLoader systemState *system.SystemState @@ -370,6 +375,14 @@ func (s *backendSupervisor) allocateFreePort(key string) (int, error) { return 0, fmt.Errorf("%w: every port tried was already in use", ErrNoFreePort) } +// backendBindHost returns the host that a backend process listens on. +func (s *backendSupervisor) backendBindHost() string { + if s.bindHost == "" { + return "0.0.0.0" + } + return s.bindHost +} + // portIsFree reports whether nothing listens on port on this host. Two checks, // because neither is enough everywhere. A connect to loopback finds a listener // bound to a specific address, which a bind of the wildcard address can miss on @@ -534,7 +547,7 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin s.mu.Unlock() return "", fmt.Errorf("allocating gRPC port for backend %s: %w", backend, err) } - bindAddr := fmt.Sprintf("0.0.0.0:%d", port) + bindAddr := net.JoinHostPort(s.backendBindHost(), strconv.Itoa(port)) clientAddr := fmt.Sprintf("127.0.0.1:%d", port) proc, err := s.ml.StartProcess(backendPath, backend, bindAddr, nil) diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 79df4e884203..8bf0401e96f1 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -4,6 +4,7 @@ import ( "cmp" "context" "encoding/json" + "errors" "fmt" "net" "os" @@ -88,58 +89,47 @@ func Run(ctx *cliContext.Context, cfg *Config) error { registrationBody := cfg.registrationBody() natsTLS := messaging.TLSFiles{CA: cfg.NatsTLSCA, Cert: cfg.NatsTLSCert, Key: cfg.NatsTLSKey} - // Resolve how to connect to NATS. Static env credentials cannot be re-minted, - // so register once and use them directly. Otherwise the credential manager - // (re)registers to obtain credentials — waiting through admin approval — and - // refreshes them before the minted JWT expires, so the connection survives - // expiry via a transparent reconnect. + // Register, and learn from the answer which carrier the cluster runs on. + // + // Static NATS credentials from the environment cannot be minted again, so + // the worker registers once and uses them as they are. Otherwise the + // credential manager registers, waits through admin approval, and for NATS + // it refreshes the credentials before the minted JWT expires, so that the + // connection survives the expiry through a reconnect that the client does by + // itself. For the tunnel it holds the token of the node. var ( - nodeID string - connectNats func() (*messaging.Client, error) - ) - if cfg.NatsJWT != "" || cfg.NatsUserSeed != "" { - nid, _, _, _, regErr := regClient.RegisterWithRetry(shutdownCtx, registrationBody, 10) - if regErr != nil { - return fmt.Errorf("failed to register with frontend: %w", regErr) - } - nodeID = nid - connectNats = func() (*messaging.Client, error) { - return connectNATS(cfg.NatsURL, cfg.NatsJWT, cfg.NatsUserSeed, "", "", cfg.NatsAuthRequired(), natsTLS) - } - } else { - credMgr := workerregistry.NewNATSCredentialManager( + res *workerregistry.RegisterResponse + credMgr = workerregistry.NewCredentialManager( func(ctx context.Context) (*workerregistry.RegisterResponse, error) { return regClient.RegisterFull(ctx, registrationBody) }, cfg.NatsAuthRequired(), ) - res, regErr := credMgr.Acquire(shutdownCtx) - if regErr != nil { - return fmt.Errorf("failed to register with frontend: %w", regErr) - } - nodeID = res.ID - connectNats = func() (*messaging.Client, error) { - var opts []messaging.Option - if credMgr.HasCredentials() { - opts = append(opts, messaging.WithUserJWTProvider(credMgr.Provider())) - } - if natsTLS.Enabled() { - opts = append(opts, messaging.WithTLS(natsTLS)) - } - client, cerr := messaging.New(cfg.NatsURL, opts...) - if cerr == nil && credMgr.HasCredentials() { - go func() { - if err := credMgr.RefreshLoop(shutdownCtx); err != nil { - xlog.Error("NATS credential refresh permanently failed; shutting down worker", "error", err) - shutdownCancel() - } - }() - } - return client, cerr - } + staticNATS = cfg.NatsJWT != "" || cfg.NatsUserSeed != "" + regErr error + ) + if staticNATS { + res, regErr = regClient.RegisterFullWithRetry(shutdownCtx, registrationBody, 10) + } else { + res, regErr = credMgr.Acquire(shutdownCtx) + } + if regErr != nil { + return fmt.Errorf("failed to register with frontend: %w", regErr) + } + nodeID := res.ID + onTunnel, err := carrierOf(res.Carrier, cfg) + if err != nil { + return err + } + // The token comes from the credential manager, which is updated by every + // registration. With static NATS credentials there is no manager, and a + // frontend that answers with the tunnel gives the token once. + tunnelToken := credMgr.TunnelToken + if staticNATS { + tunnelToken = func() string { return res.TunnelToken } } - xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cfg.RegisterTo) + xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cfg.RegisterTo, "carrier", cmp.Or(res.Carrier, "nats")) heartbeatInterval, err := time.ParseDuration(cfg.HeartbeatInterval) if err != nil && cfg.HeartbeatInterval != "" { xlog.Warn("invalid heartbeat interval, using default 10s", "input", cfg.HeartbeatInterval, "error", err) @@ -148,7 +138,8 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // Start HTTP file transfer server. (Empty-token enforcement is handled at // the top of Run so the worker fails before registering.) - httpAddr := cfg.resolveHTTPAddr() + bindHost := cfg.bindHost(onTunnel) + httpAddr := cfg.resolveHTTPAddr(bindHost) stagingDir := filepath.Join(cfg.ModelsPath, "..", "staging") cacheDir := filepath.Join(cfg.ModelsPath, "..", "cache") dataDir := filepath.Join(cfg.ModelsPath, "..", "data") @@ -181,16 +172,49 @@ func Run(ctx *cliContext.Context, cfg *Config) error { return fmt.Errorf("starting HTTP file transfer server: %w", err) } - // Connect to NATS - xlog.Info("Connecting to NATS", "url", sanitize.URL(cfg.NatsURL)) - natsClient, err := connectNats() - if err != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("connecting to NATS: %w", err) + // Attach to the carrier that the frontend named. The link is what the + // heartbeat and the readiness probe look at: a worker whose link is down + // still sends its HTTP heartbeat, and the registry would show a node that + // looks healthy and cannot be reached. + var ( + link carrierLink + natsClient *messaging.Client + linkProbe func() error + linkDownMsg string + ) + if onTunnel { + t, terr := StartTunnel(shutdownCtx, TunnelConfig{ + FrontendURL: cfg.RegisterTo, + NodeID: nodeID, + Token: tunnelToken, + // Built by tunnelServices and not inline, so that the routing table, + // which is the security boundary of the tunnel, can be reached from a + // spec without starting a worker. + Services: tunnelServices(cfg, httpAddr), + }) + if terr != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("starting the worker tunnel: %w", terr) + } + defer func() { + if err := t.Close(); err != nil { + xlog.Warn("Closing the worker tunnel failed", "error", err) + } + }() + link, linkProbe, linkDownMsg = t, nodes.TunnelReadiness(t), "tunnel disconnected" + } else { + xlog.Info("Connecting to NATS", "url", sanitize.URL(cfg.NatsURL)) + var cerr error + natsClient, cerr = connectNatsFor(shutdownCtx, shutdownCancel, cfg, credMgr, staticNATS, natsTLS) + if cerr != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("connecting to NATS: %w", cerr) + } + defer natsClient.Close() + link, linkProbe, linkDownMsg = natsLink{natsClient}, nodes.NATSReadiness(natsClient), "NATS disconnected" } - defer natsClient.Close() - // Start heartbeat goroutine (after NATS is connected so IsConnected check works) + // Start heartbeat goroutine (after the carrier is attached so the check works) go func() { ticker := time.NewTicker(heartbeatInterval) defer ticker.Stop() @@ -199,8 +223,8 @@ func Run(ctx *cliContext.Context, cfg *Config) error { case <-shutdownCtx.Done(): return case <-ticker.C: - if !natsClient.IsConnected() { - xlog.Warn("Skipping heartbeat: NATS disconnected") + if !link.Connected() { + xlog.Warn("Skipping heartbeat: " + linkDownMsg) continue } body := cfg.heartbeatBody() @@ -237,17 +261,18 @@ func Run(ctx *cliContext.Context, cfg *Config) error { nextPort: basePort, minPort: basePort, maxPort: cfg.effectiveMaxPort(basePort), + bindHost: bindHost, } // Arm the readiness gate now that the worker can actually receive work. - // NATS is already connected at this point, so a worker that is up but cut - // off from the bus reports 503 instead of a meaningless 200 (#10987). + // The carrier is already attached at this point, so a worker that is up but + // cut off from it reports 503 instead of a meaningless 200 (#10987). // - // Readiness also covers the data path: a worker whose NATS link is fine - // but whose backend processes have died is up and useless, and the - // scheduler cannot tell the difference from the bus alone. + // Readiness also covers the data path: a worker whose link is fine but whose + // backend processes have died is up and useless, and the scheduler cannot + // tell the difference from the link alone. readiness.Set(nodes.CompositeReadiness( - nodes.NATSReadiness(natsClient), + linkProbe, nodes.BackendDataPathReadiness(supervisor, func(addr string) error { conn, err := net.DialTimeout("tcp", addr, 2*time.Second) if err != nil { @@ -267,18 +292,24 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // The watchdog stops load operations the controller no longer renews. go supervisor.runOperationWatchdog(shutdownCtx) - control := newNATSControlServer(natsClient, nodeID) - if err := supervisor.registerLifecycleVerbs(control); err != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("subscribing to worker lifecycle events: %w", err) - } - - // Serve the file staging verbs only when S3 is configured - if cfg.StorageURL != "" { - if err := cfg.registerFileStagingVerbs(control, ephemeralCapacity); err != nil { + if natsClient != nil { + control := newNATSControlServer(natsClient, nodeID) + if err := supervisor.registerLifecycleVerbs(control); err != nil { nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("subscribing to file staging subjects: %w", err) + return fmt.Errorf("subscribing to worker lifecycle events: %w", err) } + + // Serve the file staging verbs only when S3 is configured + if cfg.StorageURL != "" { + if err := cfg.registerFileStagingVerbs(control, ephemeralCapacity); err != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("subscribing to file staging subjects: %w", err) + } + } + } else { + // The tunnel carries streams to the backends and to the file-transfer + // server of this worker. The lifecycle verbs have no route on it yet. + xlog.Warn("This worker is attached to the tunnel and serves no lifecycle verbs on it: backends are not installed or stopped by the frontend") } xlog.Info("Worker ready, waiting for backend.install events") @@ -300,3 +331,40 @@ func Run(ctx *cliContext.Context, cfg *Config) error { nodes.ShutdownFileTransferServer(httpServer) return runErr } + +// carrierLink is the connection of a worker to the carrier it is attached to. +type carrierLink interface { + Connected() bool +} + +// natsLink adapts a NATS client to carrierLink. +type natsLink struct{ c *messaging.Client } + +func (l natsLink) Connected() bool { return l.c.IsConnected() } + +// carrierOf maps the carrier that the frontend named to the one this worker +// attaches to, and reports whether it is the tunnel. +// +// A frontend that predates carriers names none, and the cluster runs on NATS. +// Then the worker needs a NATS URL, and without one it stops with a message that +// says so. The check is here and not in the flag, because the flag cannot know +// the carrier: for a cluster on the tunnel the URL is not needed. +func carrierOf(named string, cfg *Config) (onTunnel bool, err error) { + switch named { + case "tunnel": + if cfg.NatsURL != "" { + xlog.Info("The cluster runs on the tunnel, so the NATS URL of this worker is not used") + } + return true, nil + case "", "nats": + if cfg.NatsURL == "" { + if named == "" { + return false, errors.New("the frontend does not name a carrier, so the cluster runs on NATS, and this worker has no NATS URL: set LOCALAI_NATS_URL") + } + return false, errors.New("the cluster runs on NATS and this worker has no NATS URL: set LOCALAI_NATS_URL") + } + return false, nil + default: + return false, fmt.Errorf("the frontend names the carrier %q, which this worker does not know: upgrade the worker", named) + } +} From 42d4b2c47b1f53dac111b79412058c4c824d9899 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:55:28 +0000 Subject: [PATCH 030/162] feat(cli): let the agent worker learn the carrier from the registration The NATS URL flag of the agent worker is no longer required, because for a cluster on the tunnel nothing would need it. After the registration the agent worker checks the carrier that the frontend named. It needs a URL for NATS, and it stops with a message for the tunnel, which it cannot use yet. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/agent_worker.go | 23 ++++++++++++++++++++++- core/cli/agent_worker_subject_test.go | 17 +++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 9eb493314aca..9ee8ecc57928 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -38,7 +38,7 @@ import ( // localai agent-worker --nats-url nats://... --register-to http://localai:8080 type AgentWorkerCMD struct { // NATS (required) - NatsURL string `env:"LOCALAI_NATS_URL" required:"" help:"NATS server URL" group:"distributed"` + NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed while the cluster runs on NATS" group:"distributed"` // Registration (required) RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"` @@ -88,6 +88,23 @@ func validateAgentSubject(subject string) error { return nil } +// checkAgentCarrier checks that the carrier which the frontend named is one that +// this agent worker can use. It serves its agent runs over NATS only. The flag of +// the NATS URL cannot be required, because the carrier is known only after the +// registration, so the check is here. A frontend that names no carrier predates +// carriers and runs on NATS. +func checkAgentCarrier(carrier, natsURL string) error { + switch carrier { + case "", "nats": + if natsURL == "" { + return fmt.Errorf("the cluster runs on NATS and this agent worker has no NATS URL: set LOCALAI_NATS_URL") + } + return nil + default: + return fmt.Errorf("the cluster carrier is %q, which this agent worker cannot use: it serves agent runs over NATS only", carrier) + } +} + func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { if err := validateAgentSubject(cmd.Subject); err != nil { return err @@ -139,6 +156,10 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { nodeID := res.ID xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo) + if err := checkAgentCarrier(res.Carrier, cmd.NatsURL); err != nil { + return err + } + // Use provisioned API token if none was set if cmd.APIToken == "" { cmd.APIToken = res.APIToken diff --git a/core/cli/agent_worker_subject_test.go b/core/cli/agent_worker_subject_test.go index 3f5f14fa3344..9b4537407743 100644 --- a/core/cli/agent_worker_subject_test.go +++ b/core/cli/agent_worker_subject_test.go @@ -29,3 +29,20 @@ var _ = Describe("validateAgentSubject", func() { Expect(err.Error()).To(ContainSubstring("LOCALAI_AGENT_SUBJECT")) }) }) + +var _ = Describe("checkAgentCarrier", func() { + It("accepts NATS with a URL, named or not named", func() { + Expect(checkAgentCarrier("nats", "nats://n:4222")).To(Succeed()) + Expect(checkAgentCarrier("", "nats://n:4222")).To(Succeed()) + }) + + It("asks for the URL when the cluster runs on NATS and there is none", func() { + Expect(checkAgentCarrier("nats", "")).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) + Expect(checkAgentCarrier("", "")).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) + }) + + It("says that it cannot use the tunnel, with or without a URL", func() { + Expect(checkAgentCarrier("tunnel", "")).To(MatchError(ContainSubstring("cannot use"))) + Expect(checkAgentCarrier("tunnel", "nats://n:4222")).To(MatchError(ContainSubstring("cannot use"))) + }) +}) From bb5b5a9e7b259278c779364af26725ac47906622 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:55:41 +0000 Subject: [PATCH 031/162] docs(agents): describe the worker tunnel for contributors Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 39 +++++++++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index c6465927d8be..7d193955569e 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -8,7 +8,8 @@ dial for the gRPC and HTTP connections to a worker. Only the carrier files impor the NATS libraries (`core/services/messaging/client.go`, `tls.go`, `core/services/nodes/control_nats.go` and `pkg/natsauth`). Only the pgbus package imports `github.com/jackc/pgx`, and only `core/services/carrier/pgbus.go` imports -the pgbus package. A spec in `core/services/carrier` fails when any other +the pgbus package. Only `core/services/tunnel` imports yamux. A spec in +`core/services/carrier` fails when any other non-test file imports one of them; it reads the import lines of every Go file, so a build constraint does not hide one. A new carrier adds its library and its files to that spec. @@ -69,6 +70,42 @@ A NATS URL that points at a server which is not up does not stop the start: `messaging.New` retries on a failed connect, so a frontend can start before its broker. A URL it cannot parse does. +## The worker tunnel + +A worker that has no address that others can reach holds one outbound websocket +to a frontend. `core/services/tunnel` has the wire format (a request frame names +a tag and a target, and a reply frame always follows), the session, the registry +of the sessions that a replica holds, and `Splice`. The connect endpoint is +`GET /api/cluster/connect` (`core/http/endpoints/cluster`) and the client is +`core/services/worker/tunnel.go`. + +- Every websocket of the tunnel comes from `tunnel.NewUpgrader` or + `tunnel.NewDialer`. They set 64 KiB buffers. The 4 KiB default of gorilla makes + a transfer from the worker five times slower than the other direction. A spec + reads the source files and fails on a literal `websocket.Upgrader{` or + `websocket.Dialer{`. +- A worker has two sessions, called lanes. The inference lane carries model + calls, control and health, with the yamux windows at their defaults. The bulk + lane carries file transfers, with windows of 4 MiB and 32 MiB. The caller + chooses the lane in `Registry.Open`. A large transfer on the inference lane + delays every small call on it, because yamux queues up to one window ahead of + them. `Registry.Open` falls back to the inference lane when a node has no bulk + session. +- The inference lane owns the claim in `node_connections`. The bulk lane is held + under that claim, only on the replica that holds the inference lane, and it is + closed with it. A bulk dial that reaches another replica gets 409 and the + worker dials again. +- `Splice` passes the end of one direction on with `CloseWrite` and goes on + copying the other direction. It closes both streams at once when a direction + fails or when a stream cannot half-close. +- A node has its own tunnel credential. Registration mints it, returns it once, + and stores only its SHA-256 (`BackendNode.TunnelTokenHash`). The registration + token of the deployment never opens a tunnel. Every registration mints a new + credential, so the client reads it at every dial. +- A refusal of a stream is one of four. `tunnel.IsWorkerAnswer` is true for the + three that are evidence about a backend. `ErrStreamNotServed` says that the + worker learned nothing, and it must never count as evidence. + ## Queues `WorkKind` names the work (`WorkTask`, `WorkMCPCI`, `WorkAgentRun`). From d7063941c3008dba3acd316f86ce9c2e9bd8148b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 21:57:36 +0000 Subject: [PATCH 032/162] fix(tunnel): close joined streams when the session ends under a half-close After one direction of a spliced pair passed its end on, the other direction can sit in a read of the local backend. If the session under the other stream then ends, nothing writes to that stream again, so nothing wakes the read. The streams stay open, the worker waits for them before it connects again, and the tunnel cannot come back. Splice now watches the sessions of its streams while it waits for the second direction, and closes both streams when one of them ends. The connect endpoint also logs a bulk dial that lost the inference session at debug level, because the worker dials again and it is not an error of the replica. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/cluster/connect.go | 9 +++- core/services/tunnel/splice.go | 64 ++++++++++++++++++++++---- core/services/tunnel/splice_test.go | 33 ++++++++++++- 3 files changed, 95 insertions(+), 11 deletions(-) diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go index ac0aa383c237..2bd5231eef4f 100644 --- a/core/http/endpoints/cluster/connect.go +++ b/core/http/endpoints/cluster/connect.go @@ -155,7 +155,14 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *tunnel.Registry) echo // worker reads. attachment, err := tunnels.Attach(c.Request().Context(), nodeID, lane, sess) if err != nil { - xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "lane", lane, "error", err) + if errors.Is(err, tunnel.ErrNotOwner) { + // The inference session ended between the check before the + // upgrade and here. The worker dials again, and this is not an + // error of this replica. + xlog.Debug("worker bulk tunnel dial lost the inference session of its node", "node", nodeID) + } else { + xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "lane", lane, "error", err) + } _ = sess.Close() return nil } diff --git a/core/services/tunnel/splice.go b/core/services/tunnel/splice.go index 20945a41e560..cacde6e0e5df 100644 --- a/core/services/tunnel/splice.go +++ b/core/services/tunnel/splice.go @@ -4,6 +4,7 @@ import ( "errors" "io" "net" + "sync" "github.com/libp2p/go-yamux/v5" ) @@ -43,15 +44,29 @@ func Splice(a, b io.ReadWriteCloser) error { first := <-results if first.halfClosed { // The end was passed on. The other direction ends when the peer closes - // its side, or it fails when the session dies, and either is its own - // result and not an echo of a close by Splice. - second := <-results - closeErrA := a.Close() - closeErrB := b.Close() - if second.err != nil { - return second.err + // its side, and that is its own result and not an echo of a close by + // Splice. + // + // It can also be parked in a read of a stream that nobody will write to + // again, because the session under the other stream ended. Nothing wakes + // that read, and the streams would stay open for ever. So the end of + // either session closes both streams. + done := make(chan struct{}) + defer close(done) + select { + case second := <-results: + closeErrA := a.Close() + closeErrB := b.Close() + if second.err != nil { + return second.err + } + return firstCloseFailure(closeErrA, closeErrB) + case <-watchSessions(done, a, b): + _ = a.Close() + _ = b.Close() + <-results + return errSessionEnded } - return firstCloseFailure(closeErrA, closeErrB) } // Each stream is closed here and nowhere else. A second Close of a yamux @@ -82,6 +97,39 @@ func firstCloseFailure(errA, errB error) error { return normalizeStreamErr(errB) } +// errSessionEnded is the result of a Splice whose session ended while one +// direction was still open, so the request was cut. +var errSessionEnded = errors.New("tunnel: the session ended while a stream was half-closed") + +// sessionOf is implemented by the streams of a session. +type sessionOf interface { + Session() *yamux.Session +} + +// watchSessions returns a channel that is closed when the session of one of the +// streams ends. It stops watching when done is closed, so that the goroutines do +// not outlive the Splice for a session that goes on. A stream that does not +// belong to a session gives no signal. +func watchSessions(done <-chan struct{}, streams ...io.ReadWriteCloser) <-chan struct{} { + ended := make(chan struct{}) + var once sync.Once + for _, st := range streams { + s, ok := st.(sessionOf) + if !ok { + continue + } + gone := s.Session().CloseChan() + go func() { + select { + case <-gone: + once.Do(func() { close(ended) }) + case <-done: + } + }() + } + return ended +} + // copyResult is the end of one direction. halfClosed is true when the end was // passed on to the destination with CloseWrite. type copyResult struct { diff --git a/core/services/tunnel/splice_test.go b/core/services/tunnel/splice_test.go index f14f3166f0bf..d3e3285a68a2 100644 --- a/core/services/tunnel/splice_test.go +++ b/core/services/tunnel/splice_test.go @@ -17,8 +17,14 @@ import ( // on the worker) with the other end of that stream, and a TCP connection with // the other end of that connection. func streamAndTCP() (far net.Conn, spliceStream net.Conn, spliceTCP, backend *net.TCPConn) { + far, spliceStream, spliceTCP, backend, _, _ = streamAndTCPWithSessions() + return +} + +// streamAndTCPWithSessions is streamAndTCP that also returns the two sessions. +func streamAndTCPWithSessions() (far net.Conn, spliceStream net.Conn, spliceTCP, backend *net.TCPConn, frontend, worker *tunnel.Session) { GinkgoHelper() - frontend, worker := sessionPair(tunnel.LaneInference) + frontend, worker = sessionPair(tunnel.LaneInference) type accepted struct { c net.Conn err error @@ -38,7 +44,7 @@ func streamAndTCP() (far net.Conn, spliceStream net.Conn, spliceTCP, backend *ne Expect(tunnel.ReadStreamReply(got.c)).To(Succeed()) spliceTCP, backend = tcpPair() - return far, got.c, spliceTCP, backend + return far, got.c, spliceTCP, backend, frontend, worker } var _ = Describe("Splice", func() { @@ -116,6 +122,29 @@ var _ = Describe("Splice", func() { Eventually(done, 5*time.Second).Should(Receive(BeNil())) }) + It("closes both streams when the session ends while a direction is still open", func() { + // The frontend ends its side and the backend says nothing more. Then the + // session goes away. The read from the stream is at its end and nothing + // writes to the stream again, so the read from the backend would stay + // parked for ever without the watch on the session. + far, stream, local, backend, frontend, _ := streamAndTCPWithSessions() + done := make(chan error, 1) + go func() { done <- tunnel.Splice(stream, local) }() + + Expect(far.(interface{ CloseWrite() error }).CloseWrite()).To(Succeed()) + // The backend reads the end of the request. The half-close has passed. + _, err := io.ReadAll(backend) + Expect(err).ToNot(HaveOccurred()) + Consistently(done, 200*time.Millisecond).ShouldNot(Receive(), "the other direction is still open") + + Expect(frontend.Close()).To(Succeed()) + + Eventually(done, 5*time.Second).Should(Receive(MatchError(ContainSubstring("session ended")))) + // The backend sees its connection closed. + _, err = backend.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred()) + }) + It("closes both streams when one end cannot half-close", func() { a1, a2 := net.Pipe() b1, b2 := net.Pipe() From 989b73f9c3482efabf50c014ca33f12ea919e382 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 22:00:15 +0000 Subject: [PATCH 033/162] refactor(worker): drop an accessor of the credential manager that nothing calls Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/workerregistry/credentials.go | 10 ---------- core/cli/workerregistry/credentials_test.go | 1 - 2 files changed, 11 deletions(-) diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 408db9c88b95..6168876b60ef 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -56,7 +56,6 @@ type CredentialManager struct { jwt string seed string nodeID string - carrier string tunnelToken string } @@ -92,7 +91,6 @@ func (m *CredentialManager) store(res *RegisterResponse) { m.mu.Lock() defer m.mu.Unlock() m.nodeID = res.ID - m.carrier = res.Carrier if res.NatsJWT != "" && res.NatsUserSeed != "" { m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed } @@ -101,14 +99,6 @@ func (m *CredentialManager) store(res *RegisterResponse) { } } -// Carrier returns the carrier that the last registration named. It is empty for -// a frontend that predates the field, and that is a frontend with NATS. -func (m *CredentialManager) Carrier() string { - m.mu.RLock() - defer m.mu.RUnlock() - return m.carrier -} - // TunnelToken returns the token of the node for the tunnel, or an empty string // before the first registration that returned one. The tunnel client calls it // for every dial, because a registration mints a new token. diff --git a/core/cli/workerregistry/credentials_test.go b/core/cli/workerregistry/credentials_test.go index a1fd130a813c..f1f71fa5d3cf 100644 --- a/core/cli/workerregistry/credentials_test.go +++ b/core/cli/workerregistry/credentials_test.go @@ -130,7 +130,6 @@ var _ = Describe("CredentialManager", func() { Expect(res.Carrier).To(Equal("tunnel")) Expect(f.count()).To(Equal(1)) Expect(m.TunnelToken()).To(Equal("token-1")) - Expect(m.Carrier()).To(Equal("tunnel")) Expect(m.HasCredentials()).To(BeFalse(), "a tunnel worker holds no NATS credentials") }) From fdcc25f558252025e78f5c58f9a33e7390cd7f61 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 22:02:30 +0000 Subject: [PATCH 034/162] refactor(worker): remove the single-attempt registration that retry no longer uses RegisterWithRetry now calls RegisterFullWithRetry, and nothing else called Register. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/workerregistry/client.go | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index 6bec5ca95a01..feb5e882b3d7 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -105,16 +105,6 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a return &result, nil } -// Register sends a single registration request and returns the node ID and -// optional credentials (API token for agent workers, NATS JWT when configured). -func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) (nodeID, apiToken, natsJWT, natsSeed string, err error) { - res, err := c.RegisterFull(ctx, body) - if err != nil { - return "", "", "", "", err - } - return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil -} - // RegisterWithRetry retries registration with exponential backoff. func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) { res, err := c.RegisterFullWithRetry(ctx, body, maxRetries) From 9911052fbc5ee8a565a4ba5f16c73c506b168a0b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:12:09 +0000 Subject: [PATCH 035/162] fix(pgbus): bound the connections that publishing uses Publish ran on an unbounded pool, so each concurrent call opened its own connection. 1000 concurrent publishers on a stock server gave 99 "too many clients" errors, and the same burst here gave 138. Publish now takes a slot from a semaphore before its first statement and holds it until the last one ends. Config.MaxPublishers sets the number of slots, 16 by default. PublishContext waits for a slot with a context and returns the error of the context when it ends first. One replica now uses at most 25 connections: LISTEN, the spill readers and the publishers. The package comment and the contributor notes give the budget. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 6 ++ core/services/pgbus/bus.go | 80 +++++++++++++++- core/services/pgbus/listener.go | 9 ++ core/services/pgbus/publish_bound_test.go | 110 ++++++++++++++++++++++ 4 files changed, 201 insertions(+), 4 deletions(-) create mode 100644 core/services/pgbus/publish_bound_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 7d193955569e..cb576ab83d5c 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -193,6 +193,12 @@ pool is keyed by host and port and two workers can report the same address. - Delivery is at-most-once. Anything that must survive a gap belongs in a table. - A subscriber that is reconnecting misses messages. Do not read silence as evidence about a node. +- Database connections of pgbus: one replica uses at most 1 pinned LISTEN + connection, `pgbus.spillFetchers` (8) while it reads spilled rows, and + `pgbus.Config.MaxPublishers` (16 by default) while it publishes. `Publish` + waits for a free slot and never opens another connection, so a burst of + publishers cannot exhaust `max_connections`. Add the 25 connections of each + replica to the budget of the database. ## The cluster registry diff --git a/core/services/pgbus/bus.go b/core/services/pgbus/bus.go index e5b7937650f4..5d7cd29266bc 100644 --- a/core/services/pgbus/bus.go +++ b/core/services/pgbus/bus.go @@ -24,6 +24,20 @@ // one reader at a time, a replica lost half of a burst of 64 KiB broadcasts // above about 1,000 a second. Above the limits, broadcasts are dropped and // counted, and never queued without a bound. +// +// # Connections +// +// The database is shared, and a stock server accepts 100 connections for all of +// its clients. One replica of this carrier uses, at most: +// +// - 1 pinned connection for LISTEN; +// - spillFetchers (8) connections while it reads spilled rows; +// - Config.MaxPublishers (16 by default) connections while it publishes. +// +// That is 25 connections for each replica, and the rest of the application +// comes on top. Publish waits for a free slot when all of them are in use, and +// does not open another connection. Size max_connections for the replicas times +// this budget, or lower MaxPublishers. package pgbus import ( @@ -115,6 +129,11 @@ type Config struct { // SweepInterval is how often this carrier deletes the spilled rows that aged // out. Zero means Retention / 2. SweepInterval time.Duration + // MaxPublishers bounds how many Publish calls use the database at the same + // time on this replica. Each one holds a connection of the pool until its + // statement ends, so this bound is also the bound on the connections that + // publishing opens. Zero means DefaultMaxPublishers. + MaxPublishers int // Meter receives the counters of the carrier. Nil means the global meter. Meter metric.Meter } @@ -151,6 +170,9 @@ type Bus struct { retention time.Duration metrics *metrics + // publishSlots is the semaphore that bounds the concurrent publishes. See + // Config.MaxPublishers. + publishSlots chan struct{} cmds chan listenCmd inbound chan inbound listenerDone chan struct{} @@ -227,8 +249,14 @@ func New(ctx context.Context, cfg Config) (*Bus, error) { retention = DefaultSpillRetention } + publishers := cfg.MaxPublishers + if publishers <= 0 { + publishers = DefaultMaxPublishers + } + busCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) b := &Bus{ + publishSlots: make(chan struct{}, publishers), cfg: cfg, ctx: busCtx, cancel: cancel, @@ -306,7 +334,17 @@ func inlineNotification(subject string, data any) ([]byte, json.RawMessage, erro // Publish fans a message out to every subscriber of the subject on every // replica, this one included. +// +// It waits for a free publish slot when Config.MaxPublishers calls are already +// in progress. The wait ends when the carrier is closed. func (b *Bus) Publish(subject string, data any) error { + return b.PublishContext(b.ctx, subject, data) +} + +// PublishContext is Publish with a context for the wait for a publish slot and +// for the database statements. When the context ends first, it returns the error +// of the context and publishes nothing. +func (b *Bus) PublishContext(ctx context.Context, subject string, data any) error { channel, err := ChannelFor(subject) if err != nil { return err @@ -318,13 +356,29 @@ func (b *Bus) Publish(subject string, data any) error { if err := messaging.CheckBroadcastSize(subject, len(payload)); err != nil { return err } + + // The slot is taken before the first statement, and it is held until the + // last one ends. Without it, a burst of publishers opens a connection for + // each of them, and a stock server refuses the ones above its limit. + if err := b.acquirePublishSlot(ctx); err != nil { + return err + } + defer func() { <-b.publishSlots }() + + // A context that also ends when the carrier is closed, so a statement does + // not outlive Close. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + stop := context.AfterFunc(b.ctx, cancel) + defer stop() + // This is the one size decision in the package. Several subjects exceed the // limit in normal operation (the result of a job carries the whole output of // a model, and a gallery progress event carries one entry for each node), so // the spill is the usual path for them and not an error. path := pathInline if len(encoded) >= maxNotifyPayloadBytes { - id, err := b.spill(subject, payload) + id, err := b.spill(ctx, subject, payload) if err != nil { return err } @@ -337,13 +391,31 @@ func (b *Bus) Publish(subject string, data any) error { // pg_notify and not a NOTIFY statement, because the channel is a value here // and NOTIFY would need it as an identifier in the text. - if err := b.cfg.DB.WithContext(b.ctx).Exec("SELECT pg_notify(?, ?)", channel, string(encoded)).Error; err != nil { + if err := b.cfg.DB.WithContext(ctx).Exec("SELECT pg_notify(?, ?)", channel, string(encoded)).Error; err != nil { return fmt.Errorf("pgbus: publishing on %q: %w", subject, err) } b.metrics.publish(path) return nil } +// acquirePublishSlot takes one slot of the publish semaphore, or returns the +// error of the context that ended first. +func (b *Bus) acquirePublishSlot(ctx context.Context) error { + select { + case b.publishSlots <- struct{}{}: + return nil + default: + } + select { + case b.publishSlots <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-b.ctx.Done(): + return errors.New("pgbus: the carrier is closed") + } +} + // Subscribe registers a handler for every subject that matches the filter. The // grammar of the filter is the shared one: an exact subject, or `*` for a whole // token. @@ -614,9 +686,9 @@ func (s *subscription) Unsubscribe() error { // spill writes a broadcast that does not fit in a notification to a row, and // returns the id that the notification carries in its place. -func (b *Bus) spill(subject string, payload []byte) (string, error) { +func (b *Bus) spill(ctx context.Context, subject string, payload []byte) (string, error) { row := BusMessage{ID: uuid.New().String(), Subject: subject, Payload: payload} - if err := b.cfg.DB.WithContext(b.ctx).Create(&row).Error; err != nil { + if err := b.cfg.DB.WithContext(ctx).Create(&row).Error; err != nil { return "", fmt.Errorf("pgbus: spilling a broadcast on %q: %w", subject, err) } return row.ID, nil diff --git a/core/services/pgbus/listener.go b/core/services/pgbus/listener.go index 62c800578d72..890b86f12bfd 100644 --- a/core/services/pgbus/listener.go +++ b/core/services/pgbus/listener.go @@ -29,6 +29,15 @@ const DefaultQueueDepth = 1024 // deployment does not collect the outputs of models in this table. const DefaultSpillRetention = 10 * time.Minute +// DefaultMaxPublishers is how many Publish calls one replica runs against the +// database at the same time. Each one holds a pooled connection until its +// statement ends. +// +// It is twice spillFetchers: publishing is one short statement, and a replica +// that reads spilled rows and publishes at the same time stays within 25 +// connections, with the one for LISTEN. +const DefaultMaxPublishers = 2 * spillFetchers + // listenPollInterval bounds how long a LISTEN or UNLISTEN waits for the listener // goroutine to leave WaitForNotification and run it. // diff --git a/core/services/pgbus/publish_bound_test.go b/core/services/pgbus/publish_bound_test.go new file mode 100644 index 000000000000..360b47561394 --- /dev/null +++ b/core/services/pgbus/publish_bound_test.go @@ -0,0 +1,110 @@ +package pgbus_test + +import ( + "context" + "strings" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// A publish takes one connection of the pool for as long as its statement runs. +// The pool is unbounded, so a burst of publishers opens as many connections as +// there are publishers, and a stock server refuses the surplus. +var _ = Describe("bounding the connections that publishing uses", func() { + It("publishes 1000 concurrent broadcasts without a connection error", func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + newBus := func() *pgbus.Bus { + b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, Queue: 2048, Retention: time.Hour}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(b.Close) + return b + } + pub, sub := newBus(), newBus() + + var got atomic.Int64 + s, err := sub.Subscribe("jobs.bound.stream", func([]byte) { got.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + + const total = 1000 + var failures atomic.Int64 + var firstErr atomic.Value + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < total; i++ { + wg.Add(1) + go func(i int) { + defer GinkgoRecover() + defer wg.Done() + fill := "x" + if i%10 == 0 { + fill = strings.Repeat("y", 9*1024) // spills + } + <-start + if err := pub.Publish("jobs.bound.stream", map[string]string{"f": fill}); err != nil { + failures.Add(1) + firstErr.CompareAndSwap(nil, err.Error()) + } + }(i) + } + close(start) + wg.Wait() + + Expect(failures.Load()).To(BeZero(), "first error: %v", firstErr.Load()) + // Every broadcast is delivered or is in a drop counter. Nothing may go + // missing without a count. + accounted := func() int64 { + return got.Load() + int64(sub.Dropped()) + } + Eventually(accounted, 60*time.Second).Should(Equal(int64(total))) + _ = s + }) + + It("returns the error of the context when it waits for a slot that does not free up", func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + bus, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, MaxPublishers: 1}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(bus.Close) + + // A lock on the spill table holds the insert of a spilled publish, and + // with it the only slot. + held := db.Begin() + Expect(held.Error).ToNot(HaveOccurred()) + released := false + release := func() { + if !released { + released = true + Expect(held.Rollback().Error).ToNot(HaveOccurred()) + } + } + DeferCleanup(release) + Expect(held.Exec("LOCK TABLE bus_messages IN ACCESS EXCLUSIVE MODE").Error).To(Succeed()) + + first := make(chan error, 1) + go func() { + first <- bus.Publish("jobs.slot.stream", map[string]string{"f": strings.Repeat("y", 9*1024)}) + }() + Eventually(func() int64 { + var n int64 + Expect(db.Raw(`SELECT count(*) FROM pg_stat_activity + WHERE wait_event_type = 'Lock' AND query LIKE '%INSERT INTO "bus_messages"%'`).Scan(&n).Error).To(Succeed()) + return n + }, 10*time.Second, 50*time.Millisecond).Should(Equal(int64(1))) + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + err = bus.PublishContext(ctx, "jobs.slot.stream", map[string]string{"f": "x"}) + Expect(err).To(MatchError(context.DeadlineExceeded)) + + release() + Eventually(first, 10*time.Second).Should(Receive(BeNil())) + }) +}) From 08d2dabfd178a3f4e7da578b7ec8f0388857907d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:14:30 +0000 Subject: [PATCH 036/162] fix(pgbus): drop a spilled broadcast whose read does not end The dispatcher delivers in the order of arrival and waits for each spilled row to be read. A SELECT with no bound held back the delivery on every subject until it returned. The read now ends after Config.FetchTimeout, 8 seconds by default. The broadcast is dropped and counted with stage=resolve, and the ones behind it are delivered. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 3 ++ core/services/pgbus/bus.go | 17 +++++++- core/services/pgbus/fetch_timeout_test.go | 52 +++++++++++++++++++++++ core/services/pgbus/listener.go | 8 ++++ 4 files changed, 79 insertions(+), 1 deletion(-) create mode 100644 core/services/pgbus/fetch_timeout_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index cb576ab83d5c..3217d7f48eb5 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -199,6 +199,9 @@ pool is keyed by host and port and two workers can report the same address. waits for a free slot and never opens another connection, so a burst of publishers cannot exhaust `max_connections`. Add the 25 connections of each replica to the budget of the database. +- A read of a spilled row ends after `pgbus.Config.FetchTimeout` (8 seconds by + default). The broadcast is dropped and counted with `stage=resolve`, so that + one stuck read does not hold back the delivery on every subject. ## The cluster registry diff --git a/core/services/pgbus/bus.go b/core/services/pgbus/bus.go index 5d7cd29266bc..737d86524b18 100644 --- a/core/services/pgbus/bus.go +++ b/core/services/pgbus/bus.go @@ -134,6 +134,10 @@ type Config struct { // statement ends, so this bound is also the bound on the connections that // publishing opens. Zero means DefaultMaxPublishers. MaxPublishers int + // FetchTimeout bounds the read of one spilled row. A read that takes longer + // is dropped and counted, so that it does not hold back the delivery of the + // broadcasts behind it. Zero means DefaultFetchTimeout. + FetchTimeout time.Duration // Meter receives the counters of the carrier. Nil means the global meter. Meter metric.Meter } @@ -173,6 +177,7 @@ type Bus struct { // publishSlots is the semaphore that bounds the concurrent publishes. See // Config.MaxPublishers. publishSlots chan struct{} + fetchTimeout time.Duration cmds chan listenCmd inbound chan inbound listenerDone chan struct{} @@ -254,8 +259,14 @@ func New(ctx context.Context, cfg Config) (*Bus, error) { publishers = DefaultMaxPublishers } + fetchTimeout := cfg.FetchTimeout + if fetchTimeout <= 0 { + fetchTimeout = DefaultFetchTimeout + } + busCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) b := &Bus{ + fetchTimeout: fetchTimeout, publishSlots: make(chan struct{}, publishers), cfg: cfg, ctx: busCtx, @@ -694,9 +705,13 @@ func (b *Bus) spill(ctx context.Context, subject string, payload []byte) (string return row.ID, nil } +// resolveSpill reads the row of a spilled broadcast. The read ends after +// fetchTimeout, because the dispatcher waits for it in the order of arrival. func (b *Bus) resolveSpill(id string) ([]byte, error) { + ctx, cancel := context.WithTimeout(b.ctx, b.fetchTimeout) + defer cancel() var row BusMessage - if err := b.cfg.DB.WithContext(b.ctx).Where("id = ?", id).First(&row).Error; err != nil { + if err := b.cfg.DB.WithContext(ctx).Where("id = ?", id).First(&row).Error; err != nil { return nil, err } return row.Payload, nil diff --git a/core/services/pgbus/fetch_timeout_test.go b/core/services/pgbus/fetch_timeout_test.go new file mode 100644 index 000000000000..94627c58364c --- /dev/null +++ b/core/services/pgbus/fetch_timeout_test.go @@ -0,0 +1,52 @@ +package pgbus_test + +import ( + "context" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// The dispatcher takes notifications in the order they arrived, and a spilled +// one is ready only when its row is read. A read that never ends would hold +// back every subject. +var _ = Describe("a read of a spilled row that does not end", func() { + It("is dropped and counted, and the broadcasts behind it are delivered", func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + bus, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, FetchTimeout: 500 * time.Millisecond}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(bus.Close) + + other := make(chan []byte, 4) + _, err = bus.Subscribe("cache.evict", func(data []byte) { other <- data }) + Expect(err).ToNot(HaveOccurred()) + _, err = bus.Subscribe("jobs.stuck.stream", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + + // The row exists, and a lock stops every SELECT on the table. + Expect(db.Create(&pgbus.BusMessage{ID: "stuck-1", Subject: "jobs.stuck.stream", Payload: []byte(`{"i":1}`)}).Error).To(Succeed()) + held := db.Begin() + Expect(held.Error).ToNot(HaveOccurred()) + released := false + release := func() { + if !released { + released = true + Expect(held.Rollback().Error).ToNot(HaveOccurred()) + } + } + DeferCleanup(release) + Expect(held.Exec("LOCK TABLE bus_messages IN ACCESS EXCLUSIVE MODE").Error).To(Succeed()) + + Expect(db.Exec("SELECT pg_notify('localai_jobs', ?)", `{"s":"jobs.stuck.stream","i":"stuck-1"}`).Error).To(Succeed()) + Expect(bus.Publish("cache.evict", map[string]string{"model": "m"})).To(Succeed()) + + Eventually(other, 5*time.Second).Should(Receive(MatchJSON(`{"model":"m"}`)), + "one stuck read held back the delivery on every subject") + Expect(bus.Dropped()).To(Equal(uint64(1))) + }) +}) diff --git a/core/services/pgbus/listener.go b/core/services/pgbus/listener.go index 890b86f12bfd..f656401d85bf 100644 --- a/core/services/pgbus/listener.go +++ b/core/services/pgbus/listener.go @@ -38,6 +38,14 @@ const DefaultSpillRetention = 10 * time.Minute // connections, with the one for LISTEN. const DefaultMaxPublishers = 2 * spillFetchers +// DefaultFetchTimeout bounds the read of one spilled row. The read is one SELECT +// on the primary key, which ends in a millisecond when the database answers, so +// a read that takes seconds is a database that does not answer. +// +// The dispatcher waits for each notification in the order of arrival, so a read +// that has no bound holds back the delivery on every subject. +const DefaultFetchTimeout = 8 * time.Second + // listenPollInterval bounds how long a LISTEN or UNLISTEN waits for the listener // goroutine to leave WaitForNotification and run it. // From d891449b5414a7517cf51de84104a5d6eecfbfb0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:17:44 +0000 Subject: [PATCH 037/162] fix(cluster): bound the database calls of the membership loop The heartbeat, the sweep and the purge ran with the context of the application and no timeout. A statement that did not return stopped the heartbeat, and the peers then reaped a replica that still served. Stop also waited for the loop with no bound, so it could outlast the deregistration timeout. Each call of a pass now ends after one interval. The loop cancels the pass in progress when Stop is called, and Stop waits for the loop for no longer than the deregistration timeout. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/export_test.go | 16 +++++ core/services/cluster/membership.go | 60 ++++++++++++++++--- core/services/cluster/membership_test.go | 75 ++++++++++++++++++++++++ 3 files changed, 143 insertions(+), 8 deletions(-) create mode 100644 core/services/cluster/export_test.go diff --git a/core/services/cluster/export_test.go b/core/services/cluster/export_test.go new file mode 100644 index 000000000000..4728e943b7c1 --- /dev/null +++ b/core/services/cluster/export_test.go @@ -0,0 +1,16 @@ +package cluster + +import ( + "context" + "time" +) + +// These expose internals to the specs. The production code has no caller for +// them. + +// SetInterval changes how often the loop ticks, and how long one database call +// of a tick may take. Call it before Start. +func (m *Membership) SetInterval(d time.Duration) { m.interval = d } + +// Tick runs one pass of the loop on the calling goroutine. +func (m *Membership) Tick(ctx context.Context) { m.tick(ctx) } diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index 6ea2662843b1..59e5df771ffa 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -208,7 +208,16 @@ func (m *Membership) Stop() { m.stopOnce.Do(func() { close(m.stop) }) // Only a started Membership closes done. Waiting on one that never // started would block for ever. - <-m.done + // + // The loop cancels the call in progress when stop closes, so this ends + // at once. The bound is for a driver that does not honour the + // cancellation: shutdown must not wait for a database for longer than + // the deregistration does. + select { + case <-m.done: + case <-time.After(deregisterTimeout): + xlog.Warn("The cluster membership loop did not stop in time; deregistering anyway", "id", m.id, "within", deregisterTimeout) + } } // This is not the context that Start received: that is the context of the @@ -226,6 +235,18 @@ func (m *Membership) Stop() { func (m *Membership) loop(ctx context.Context) { defer close(m.done) + // The calls of a pass end when Stop is called, and not only when the + // application ends. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + go func() { + select { + case <-m.stop: + cancel() + case <-ctx.Done(): + } + }() + ticker := time.NewTicker(m.interval) defer ticker.Stop() @@ -241,19 +262,32 @@ func (m *Membership) loop(ctx context.Context) { } } +// call returns the context of one database call of a pass. A call that takes +// longer than the interval is cancelled. Without the bound, a statement that +// does not return stops the heartbeat, and the peers then reap this replica +// while it still serves. +func (m *Membership) call(ctx context.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(ctx, m.interval) +} + // tick refreshes the row of this replica and sweeps the dead ones. // // Every replica sweeps and there is no elected sweeper. The deletes are // idempotent and cheap, and an elected sweeper is one more process that must be // alive for the cluster to notice that another is not. func (m *Membership) tick(ctx context.Context) { - err := m.reg.Heartbeat(ctx, m.id) + hbCtx, cancel := m.call(ctx) + err := m.reg.Heartbeat(hbCtx, m.id) + cancel() if errors.Is(err, ErrInstanceNotFound) { // Another replica swept this row while this process stalled for long // enough to look dead. Register again and do not only heartbeat: the // row has to be rebuilt. xlog.Warn("Cluster instance row was reaped, registering again", "id", m.id) - if err := m.register(ctx); err != nil { + regCtx, cancel := m.call(ctx) + err := m.register(regCtx) + cancel() + if err != nil { // The connections are not claimed again in this case. A claim // written now would name an instance row that does not exist, and // the next sweep would delete it as an orphan. The sweep below @@ -262,13 +296,17 @@ func (m *Membership) tick(ctx context.Context) { } else { m.reclaim(ctx) } - } else if err != nil { + } else if err != nil && ctx.Err() == nil { xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) } - instances, cleared, err := m.reg.ReapStale(ctx, m.id, m.liveness) + reapCtx, cancel := m.call(ctx) + instances, cleared, err := m.reg.ReapStale(reapCtx, m.id, m.liveness) + cancel() if err != nil { - xlog.Warn("Reaping stale cluster instances failed", "error", err) + if ctx.Err() == nil { + xlog.Warn("Reaping stale cluster instances failed", "error", err) + } return } if instances > 0 || cleared > 0 { @@ -283,9 +321,13 @@ func (m *Membership) tick(ctx context.Context) { m.mu.Lock() retention := m.retention m.mu.Unlock() - purged, err := m.reg.PurgeDepartedBefore(ctx, retention) + purgeCtx, cancel := m.call(ctx) + purged, err := m.reg.PurgeDepartedBefore(purgeCtx, retention) + cancel() if err != nil { - xlog.Warn("Purging departed worker connections failed", "error", err) + if ctx.Err() == nil { + xlog.Warn("Purging departed worker connections failed", "error", err) + } return } if purged > 0 { @@ -303,6 +345,8 @@ func (m *Membership) reclaim(ctx context.Context) { if reclaimer == nil { return } + ctx, cancel := m.call(ctx) + defer cancel() n, err := reclaimer.Reclaim(ctx) if err != nil { // The loop goes on. The next heartbeat fails in the same way while the diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go index 6c4c9de0f153..ebc57060c192 100644 --- a/core/services/cluster/membership_test.go +++ b/core/services/cluster/membership_test.go @@ -495,3 +495,78 @@ var _ = Describe("Claiming connections again after a sweep", func() { Eventually(instanceRows("me"), 3*cluster.InstanceHeartbeat, 250*time.Millisecond).Should(Equal(int64(1))) }) }) + +var _ = Describe("A database call of the membership loop that does not return", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + // lockInstances stops every statement on the instances table until the + // returned function runs. + lockInstances := func() func() { + GinkgoHelper() + held := db.Begin() + Expect(held.Error).ToNot(HaveOccurred()) + released := false + release := func() { + if !released { + released = true + Expect(held.Rollback().Error).ToNot(HaveOccurred()) + } + } + DeferCleanup(release) + Expect(held.Exec("LOCK TABLE instances IN ACCESS EXCLUSIVE MODE").Error).To(Succeed()) + return release + } + + It("ends the pass within the interval of each call, so a hung statement does not stall the loop for good", func() { + membership := cluster.NewMembership(reg, "me", "v1") + membership.SetInterval(300 * time.Millisecond) + Expect(reg.Register(ctx, "me", "v1", 0, "")).To(Succeed()) + lockInstances() + + done := make(chan struct{}) + go func() { + defer GinkgoRecover() + defer close(done) + membership.Tick(ctx) + }() + Eventually(done, 5*time.Second).Should(BeClosed(), + "a pass with a hung statement never ended, so the heartbeat of this replica stopped") + }) + + It("lets Stop return within its bound", func() { + membership := cluster.NewMembership(reg, "me", "v1") + membership.SetInterval(200 * time.Millisecond) + Expect(membership.Start(ctx)).To(Succeed()) + lockInstances() + + // Wait until a statement of the loop waits for the lock, so that Stop + // meets a pass that is in progress. + Eventually(func() int64 { + var n int64 + Expect(db.Raw(`SELECT count(*) FROM pg_stat_activity + WHERE wait_event_type = 'Lock' AND query LIKE '%instances%'`).Scan(&n).Error).To(Succeed()) + return n + }, 5*time.Second, 50*time.Millisecond).Should(BeNumerically(">", 0)) + + stopped := make(chan struct{}) + go func() { + defer GinkgoRecover() + defer close(stopped) + membership.Stop() + }() + // The bound of the deregistration is 5 seconds, and the wait for the loop + // must not add to it. + Eventually(stopped, 8*time.Second).Should(BeClosed(), "Stop waited for a statement that does not return") + }) +}) From d5f58a129a59139ed73eb95a9a0823f7cd6d7a86 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:21:10 +0000 Subject: [PATCH 038/162] fix(carrier): refuse a subscription to a control root in the holder The broadcaster holder attached every subscription to every set it listened on. A carrier with no request and reply, such as pgbus, refuses the control roots, so a subscription to nodes or mcp that the holder accepted would have failed the Listen of a carrier switch later. Control traffic uses its own clients and never goes through the holder. Subscribe now refuses a subject outside the broadcast roots with ErrUnservedSubject, and a subscription that it accepts is servable by every carrier. The specs that used a nodes subject on the holder now use a broadcast root. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 6 + core/application/distributed_carrier_test.go | 4 +- core/services/carrier/broadcaster.go | 12 ++ core/services/carrier/broadcaster_test.go | 130 ++++++++++++------- core/services/carrier/set.go | 3 + 5 files changed, 109 insertions(+), 46 deletions(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 3217d7f48eb5..7357a053393c 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -188,6 +188,12 @@ pool is keyed by host and port and two workers can report the same address. pgbus, uses `messaging.ValidateBroadcastSubject`, which refuses the control roots with the same error. `messagingtest.BroadcastRootSubjects` has one subject for each broadcast root, and a spec keeps it equal to the rules. +- The `carrier.Broadcaster` holder carries fan-out only. Its `Subscribe` + refuses a subject outside the broadcast roots with + `messaging.ErrUnservedSubject`, so a subscription that it accepted is + servable by every carrier and `Listen` never fails for the subject. Control + traffic (the `nodes` and `mcp` roots) does not go through the holder: it uses + the request and reply clients of the set. - Wildcards: only a whole single token `*`, never the root. `>` is refused with `messaging.ErrUnsupportedWildcard`. - Delivery is at-most-once. Anything that must survive a gap belongs in a table. diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go index ad187afed869..76fdae9f601e 100644 --- a/core/application/distributed_carrier_test.go +++ b/core/application/distributed_carrier_test.go @@ -116,10 +116,10 @@ var _ = Describe("distributed startup and the cluster carrier", func() { // A message published through the holder reaches a subscriber made // through the same holder, over the real server. got := make(chan []byte, 1) - sub, err := svc.Broadcaster.Subscribe("nodes.carrier-spec", func(b []byte) { got <- b }) + sub, err := svc.Broadcaster.Subscribe("jobs.carrier-spec", func(b []byte) { got <- b }) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = sub.Unsubscribe() }) - Expect(svc.Broadcaster.Publish("nodes.carrier-spec", "hello")).To(Succeed()) + Expect(svc.Broadcaster.Publish("jobs.carrier-spec", "hello")).To(Succeed()) Eventually(got).Should(Receive(Equal([]byte(`"hello"`)))) }) diff --git a/core/services/carrier/broadcaster.go b/core/services/carrier/broadcaster.go index 882059daf9a5..ebf8702a6731 100644 --- a/core/services/carrier/broadcaster.go +++ b/core/services/carrier/broadcaster.go @@ -59,7 +59,19 @@ func (b *Broadcaster) Publish(subject string, data any) error { // Subscribe attaches handler to every set the holder listens on. It fails, and // leaves nothing attached, when any of them refuses. +// +// The holder carries fan-out and nothing else. It refuses, with +// messaging.ErrUnservedSubject, a subject outside the broadcast roots: the +// control roots (request and reply to one node or one agent worker) have their +// own clients, and a carrier with no request and reply, such as pgbus, cannot +// serve them. A subscription that the holder accepted here and the target set +// refused later would fail the Listen of a carrier switch, long after the call +// that caused it. So a subscription that is accepted is servable by every +// carrier, and Listen does not fail for the subject. func (b *Broadcaster) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { + if err := messaging.ValidateBroadcastSubject(subject); err != nil { + return nil, err + } sub := &subscription{owner: b, subject: subject, handler: handler, inner: map[*Set]messaging.Subscription{}} // Read the sets and register in one step, so a concurrent Listen either diff --git a/core/services/carrier/broadcaster_test.go b/core/services/carrier/broadcaster_test.go index a1160bb42079..b2f6950de8c0 100644 --- a/core/services/carrier/broadcaster_test.go +++ b/core/services/carrier/broadcaster_test.go @@ -51,83 +51,83 @@ var _ = Describe("Broadcaster holder", func() { It("publishes on the carrier the pointer names", func() { got := &received{} - _, err := a.bus.Subscribe("nodes.x", got.handler()) + _, err := a.bus.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) - _, err = b.bus.Subscribe("nodes.x", got.handler()) + _, err = b.bus.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) - Expect(h.Publish("nodes.x", "one")).To(Succeed()) + Expect(h.Publish("jobs.x", "one")).To(Succeed()) Expect(got.all()).To(HaveLen(1)) - Expect(a.bus.PublishCount("nodes.x")).To(Equal(1)) - Expect(b.bus.PublishCount("nodes.x")).To(Equal(0)) + Expect(a.bus.PublishCount("jobs.x")).To(Equal(1)) + Expect(b.bus.PublishCount("jobs.x")).To(Equal(0)) cur.Store(b.set) - Expect(h.Publish("nodes.x", "two")).To(Succeed()) - Expect(a.bus.PublishCount("nodes.x")).To(Equal(1)) - Expect(b.bus.PublishCount("nodes.x")).To(Equal(1)) + Expect(h.Publish("jobs.x", "two")).To(Succeed()) + Expect(a.bus.PublishCount("jobs.x")).To(Equal(1)) + Expect(b.bus.PublishCount("jobs.x")).To(Equal(1)) }) It("delivers a subscription that was made before and after a swap", func() { got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) Expect(h.Listen(b.set)).To(Succeed()) cur.Store(b.set) // A peer that has not flipped yet still publishes on the old carrier. - Expect(a.bus.Publish("nodes.x", "from-old")).To(Succeed()) - Expect(h.Publish("nodes.x", "from-new")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "from-old")).To(Succeed()) + Expect(h.Publish("jobs.x", "from-new")).To(Succeed()) Expect(got.all()).To(ConsistOf(`"from-old"`, `"from-new"`)) }) It("listens on both carriers while a target is attached, and publishes on one", func() { got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) Expect(h.Listen(b.set)).To(Succeed()) - Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) - Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "a")).To(Succeed()) + Expect(b.bus.Publish("jobs.x", "b")).To(Succeed()) Expect(got.all()).To(ConsistOf(`"a"`, `"b"`)) - Expect(h.Publish("nodes.x", "c")).To(Succeed()) - Expect(a.bus.PublishCount("nodes.x")).To(Equal(2), "the active carrier carries the publish") - Expect(b.bus.PublishCount("nodes.x")).To(Equal(1), "the target carrier carries none") + Expect(h.Publish("jobs.x", "c")).To(Succeed()) + Expect(a.bus.PublishCount("jobs.x")).To(Equal(2), "the active carrier carries the publish") + Expect(b.bus.PublishCount("jobs.x")).To(Equal(1), "the target carrier carries none") }) It("subscribes on every attached carrier when a subscription is made during the window", func() { Expect(h.Listen(b.set)).To(Succeed()) got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) - Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) - Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "a")).To(Succeed()) + Expect(b.bus.Publish("jobs.x", "b")).To(Succeed()) Expect(got.all()).To(ConsistOf(`"a"`, `"b"`)) }) It("removes a subscription from every carrier when it is unsubscribed", func() { got := &received{} - sub, err := h.Subscribe("nodes.x", got.handler()) + sub, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) Expect(h.Listen(b.set)).To(Succeed()) Expect(sub.Unsubscribe()).To(Succeed()) - Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) - Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "a")).To(Succeed()) + Expect(b.bus.Publish("jobs.x", "b")).To(Succeed()) Expect(got.all()).To(BeEmpty()) // It stays gone for a carrier attached later. c := newFakeCarrier(cluster.CarrierTunnel, 3) Expect(h.Listen(c.set)).To(Succeed()) - Expect(c.bus.Publish("nodes.x", "c")).To(Succeed()) + Expect(c.bus.Publish("jobs.x", "c")).To(Succeed()) Expect(got.all()).To(BeEmpty()) }) It("tolerates a second Unsubscribe", func() { - sub, err := h.Subscribe("nodes.x", func([]byte) {}) + sub, err := h.Subscribe("jobs.x", func([]byte) {}) Expect(err).ToNot(HaveOccurred()) Expect(sub.Unsubscribe()).To(Succeed()) Expect(sub.Unsubscribe()).To(Succeed()) @@ -135,14 +135,14 @@ var _ = Describe("Broadcaster holder", func() { It("drops the old carrier's subscriptions on Release and keeps the new ones", func() { got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) Expect(h.Listen(b.set)).To(Succeed()) cur.Store(b.set) Expect(h.Release(a.set)).To(Succeed()) - Expect(a.bus.Publish("nodes.x", "late-old")).To(Succeed()) - Expect(b.bus.Publish("nodes.x", "new")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "late-old")).To(Succeed()) + Expect(b.bus.Publish("jobs.x", "new")).To(Succeed()) Expect(got.all()).To(Equal([]string{`"new"`})) }) @@ -152,18 +152,18 @@ var _ = Describe("Broadcaster holder", func() { It("treats a second Listen of the same carrier as nothing to do", func() { got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) Expect(h.Listen(b.set)).To(Succeed()) Expect(h.Listen(b.set)).To(Succeed()) - Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(b.bus.Publish("jobs.x", "b")).To(Succeed()) Expect(got.all()).To(HaveLen(1), "one subscription per carrier, not two") }) It("leaves nothing behind when a carrier cannot take a subscription", func() { got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).ToNot(HaveOccurred()) broken := newFakeCarrier(cluster.CarrierTunnel, 3) @@ -174,7 +174,7 @@ var _ = Describe("Broadcaster holder", func() { Expect(err).To(MatchError(boom)) // The target is not attached, and a later subscription does not try it. - _, err = h.Subscribe("nodes.y", func([]byte) {}) + _, err = h.Subscribe("jobs.y", func([]byte) {}) Expect(err).ToNot(HaveOccurred()) Expect(h.Release(broken.set)).To(Succeed(), "releasing an unattached carrier is harmless") }) @@ -186,10 +186,10 @@ var _ = Describe("Broadcaster holder", func() { Expect(h.Listen(refusing.set)).To(Succeed(), "nothing is registered yet, so nothing is refused") got := &received{} - _, err := h.Subscribe("nodes.x", got.handler()) + _, err := h.Subscribe("jobs.x", got.handler()) Expect(err).To(MatchError(boom)) - Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "a")).To(Succeed()) Expect(got.all()).To(BeEmpty(), "the half-made subscription was removed") }) @@ -247,7 +247,7 @@ var _ = Describe("Broadcaster holder", func() { entered := make(chan struct{}, 1) slow.set.Broadcaster = blockingBroadcaster{Broadcaster: slow.bus, gate: gate, entered: entered} - _, err := h.Subscribe("nodes.x", func([]byte) {}) + _, err := h.Subscribe("jobs.x", func([]byte) {}) Expect(err).ToNot(HaveOccurred()) listenDone := make(chan error, 1) @@ -255,12 +255,12 @@ var _ = Describe("Broadcaster holder", func() { Eventually(entered).Should(Receive()) published := make(chan error, 1) - go func() { published <- h.Publish("nodes.x", "during") }() + go func() { published <- h.Publish("jobs.x", "during") }() Eventually(published, time.Second).Should(Receive(BeNil())) subscribed := make(chan error, 1) go func() { - _, err := h.Subscribe("nodes.z", func([]byte) {}) + _, err := h.Subscribe("jobs.z", func([]byte) {}) subscribed <- err }() Consistently(listenDone, 50*time.Millisecond).ShouldNot(Receive()) @@ -278,12 +278,12 @@ var _ = Describe("Broadcaster holder", func() { cur.Store(stuck.set) first := make(chan error, 1) - go func() { first <- h.Publish("nodes.x", "old") }() + go func() { first <- h.Publish("jobs.x", "old") }() Eventually(entered).Should(Receive()) cur.Store(b.set) - Expect(h.Publish("nodes.x", "new")).To(Succeed()) - Expect(b.bus.PublishCount("nodes.x")).To(Equal(1)) + Expect(h.Publish("jobs.x", "new")).To(Succeed()) + Expect(b.bus.PublishCount("jobs.x")).To(Equal(1)) close(gate) Eventually(first, time.Second).Should(Receive(BeNil())) @@ -302,7 +302,7 @@ var _ = Describe("Broadcaster holder", func() { go func() { defer wg.Done() <-start - _, err := h.Subscribe("nodes.x", counts[i].handler()) + _, err := h.Subscribe("jobs.x", counts[i].handler()) Expect(err).ToNot(HaveOccurred()) }() } @@ -315,8 +315,8 @@ var _ = Describe("Broadcaster holder", func() { close(start) wg.Wait() - Expect(a.bus.Publish("nodes.x", "a")).To(Succeed()) - Expect(b.bus.Publish("nodes.x", "b")).To(Succeed()) + Expect(a.bus.Publish("jobs.x", "a")).To(Succeed()) + Expect(b.bus.Publish("jobs.x", "b")).To(Succeed()) for i := range counts { Expect(counts[i].all()).To(ConsistOf(`"a"`, `"b"`), "subscription %d", i) } @@ -334,7 +334,7 @@ var _ = Describe("Broadcaster holder", func() { var payload any = "payload" allocs := testing.AllocsPerRun(1000, func() { - _ = holder.Publish("nodes.x", payload) + _ = holder.Publish("jobs.x", payload) }) Expect(allocs).To(BeZero()) Expect(quiet.published.Load()).To(BeNumerically(">", 1000)) @@ -393,3 +393,45 @@ func (q *quietBroadcaster) Publish(string, any) error { q.published.Add(1); retu func (q *quietBroadcaster) Subscribe(string, func([]byte)) (messaging.Subscription, error) { return nil, nil } + +var _ = Describe("Broadcaster holder subjects", func() { + var ( + a, b *fakeCarrier + cur atomic.Pointer[carrier.Set] + h *carrier.Broadcaster + ) + + BeforeEach(func() { + a = newFakeCarrier(cluster.CarrierNATS, 1) + b = newFakeCarrier(cluster.CarrierTunnel, 2) + cur = atomic.Pointer[carrier.Set]{} + cur.Store(a.set) + h = carrier.NewBroadcaster(&cur) + }) + + // Control traffic has its own clients and never goes through the holder. A + // carrier with no request and reply refuses these roots, so a subscription + // that the holder accepted would fail the first Listen onto that carrier, + // long after the call that caused it. + DescribeTable("refuses a subscription to a control root when it is made", + func(subject string) { + got := &received{} + sub, err := h.Subscribe(subject, got.handler()) + Expect(err).To(MatchError(messaging.ErrUnservedSubject)) + Expect(sub).To(BeNil()) + _ = a.bus.Publish(subject, "x") + Expect(got.all()).To(BeEmpty(), "the refused subscription was attached anyway") + + Expect(h.Listen(b.set)).To(Succeed(), "a refused subscription must not stay registered for a later swap") + }, + Entry("nodes", "nodes.x.backend.install"), + Entry("mcp", "mcp.tools.execute"), + Entry("a subject outside every root", "elsewhere.x"), + ) + + It("accepts a subscription to a broadcast root", func() { + sub, err := h.Subscribe("jobs.x.progress", func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + Expect(sub).ToNot(BeNil()) + }) +}) diff --git a/core/services/carrier/set.go b/core/services/carrier/set.go index 2a5cade12124..f98e7c5e9df7 100644 --- a/core/services/carrier/set.go +++ b/core/services/carrier/set.go @@ -38,6 +38,9 @@ type Set struct { // Epoch is the epoch of the cluster row the set was built for. Epoch int64 + // Broadcaster carries fan-out and nothing else. A subject outside the + // broadcast roots (see messaging.ValidateBroadcastSubject) is not served + // through it: control traffic uses Commands, Clients and Agents. Broadcaster messaging.Broadcaster // OnReconnect registers a callback the carrier runs after it recovers a // lost connection. It is nil for a carrier with no such event. From acc225a0d59530ea4b78a4a82beb1f97e24e3740 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:21:42 +0000 Subject: [PATCH 039/162] docs(distributed): state the 8 MiB limit of a broadcast Every carrier refuses a broadcast payload above 8 MiB with ErrPayloadTooLarge. Only a NATS server that sets max_payload above 8 MiB sees a change, because the default of the server is 1 MiB. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- docs/content/features/distributed-mode.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 2fe4d4f7f16f..44c091a0163c 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -206,6 +206,12 @@ The call is idempotent. A repeat retries the stop and never extends the hold. A Upgrade the frontends first. A worker that predates operations ignores the new request fields and does not report `reports_operations`. The frontend then treats the node as legacy: it cannot confirm a stop, so a failed load holds the model for the 45 minute load deadline, as it did before leases existed, and never longer. For such a node the stop, including a cancel, is sent by exact process address, never by model name. If the address is not known, no stop is claimed, and the model is held for the 45 minutes. A new worker that gets an install from an older frontend tracks it as an anonymous operation: it kills it at its deadline only, never for missing renewals. +### Broadcast size limit + +A broadcast is a message that goes to every frontend replica, such as a job progress event or a gallery update. Every carrier now refuses a broadcast with a payload above 8 MiB. The call fails with `ErrPayloadTooLarge`, and nothing is sent. The limit is the same for every carrier, so a deployment keeps its behavior when the carrier changes. + +This matters only if your NATS server sets `max_payload` above 8 MiB. The default of the NATS server is 1 MiB, and a message above the server limit was always refused. Keep `max_payload` at 8 MiB or lower. + ### NATS JWT authentication (recommended for production) By default, NATS connections are anonymous: any client that can reach port `4222` may publish control-plane subjects such as `nodes..backend.install`. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential. From 1c925583de5408a5ff5d3fdea950d86202f58ee3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:22:19 +0000 Subject: [PATCH 040/162] test(cluster): fail on an ordered comparison of an epoch Epochs from Claim are unique and not ordered, so the only correct comparison is equality. The note was on Claim only. It is now on Owner and Release, and in the contributor notes. A spec parses the non-test code under core/services and fails on a <, >, <= or >= that has an operand named like an epoch. It is a syntactic check, so it cannot see a comparison through another name. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 8 +- core/services/cluster/epoch_order_test.go | 110 ++++++++++++++++++++++ core/services/cluster/ownership.go | 11 ++- 3 files changed, 126 insertions(+), 3 deletions(-) create mode 100644 core/services/cluster/epoch_order_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 7357a053393c..823f38f6d360 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -220,8 +220,12 @@ replica has built a carrier (`ReadyEpoch`) and the reason when it could not (`ReadyReason`). Table `node_connections` records which live replica holds the connection of a worker. `Registry.Claim` takes a new epoch from a sequence for every claim, and `Release` needs both the owner and the epoch, so a replica that -lost a worker cannot clear the claim of the one that won it. Compare epochs for -equality and not for order. `Registry.Presence` tells a worker that is +lost a worker cannot clear the claim of the one that won it. Epochs are unique, and +they are not ordered: a claim that inserts after a `Release` can draw a number +that is lower than one already issued. Compare epochs for equality (`==`, +`!=`) and never for order. A spec (`core/services/cluster/epoch_order_test.go`) +fails on a `<`, `>`, `<=` or `>=` with an operand named like an epoch in code +under `core/services`. `Registry.Presence` tells a worker that is reconnecting from one that is gone only when the departure is older than the grace, and nothing but `PresenceGone` may be read as absence. diff --git a/core/services/cluster/epoch_order_test.go b/core/services/cluster/epoch_order_test.go new file mode 100644 index 000000000000..62994cbbcdac --- /dev/null +++ b/core/services/cluster/epoch_order_test.go @@ -0,0 +1,110 @@ +package cluster_test + +import ( + "go/ast" + "go/parser" + "go/token" + "io/fs" + "path/filepath" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// The epoch that Registry.Claim returns is unique, and it is not ordered +// between claims. The only comparison that is correct is equality. This is a +// syntactic check: it finds an operand whose name contains "epoch" under <, >, +// <= or >=. It cannot see a comparison that goes through another name, and the +// doc comment on Claim is the rule that it backs up. + +// orderedEpochComparisons returns the positions of the ordered comparisons in +// src that have an operand named like an epoch. +func orderedEpochComparisons(filename string, src any) ([]string, error) { + fset := token.NewFileSet() + file, err := parser.ParseFile(fset, filename, src, 0) + if err != nil { + return nil, err + } + var found []string + ast.Inspect(file, func(n ast.Node) bool { + bin, ok := n.(*ast.BinaryExpr) + if !ok { + return true + } + switch bin.Op { + case token.LSS, token.GTR, token.LEQ, token.GEQ: + default: + return true + } + if namesEpoch(bin.X) || namesEpoch(bin.Y) { + found = append(found, fset.Position(bin.Pos()).String()) + } + return true + }) + return found, nil +} + +// namesEpoch reports whether the operand is an identifier or a field whose name +// contains "epoch", at any depth of the operand. +func namesEpoch(e ast.Expr) bool { + named := false + ast.Inspect(e, func(n ast.Node) bool { + switch v := n.(type) { + case *ast.Ident: + named = named || strings.Contains(strings.ToLower(v.Name), "epoch") + case *ast.SelectorExpr: + named = named || strings.Contains(strings.ToLower(v.Sel.Name), "epoch") + } + return !named + }) + return named +} + +var _ = Describe("Comparing epochs", func() { + It("finds an ordered comparison of an epoch", func() { + for _, body := range []string{ + `if conn.Epoch > held {}`, + `if mine < row.Epoch {}`, + `if epoch >= 3 {}`, + `if a.readyEpoch <= b {}`, + } { + found, err := orderedEpochComparisons("x.go", "package x\nfunc f() {\n"+body+"\n}\n") + Expect(err).ToNot(HaveOccurred()) + Expect(found).To(HaveLen(1), body) + } + }) + + It("lets equality and unrelated comparisons through", func() { + found, err := orderedEpochComparisons("x.go", `package x +func f() { + if conn.Epoch == held || a != b.Epoch {} + if n > 0 {} +} +`) + Expect(err).ToNot(HaveOccurred()) + Expect(found).To(BeEmpty()) + }) + + It("is not done anywhere in the services", func() { + root := ".." // core/services + var offenders []string + err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() || !strings.HasSuffix(path, ".go") || strings.HasSuffix(path, "_test.go") { + return nil + } + found, err := orderedEpochComparisons(path, nil) + if err != nil { + return err + } + offenders = append(offenders, found...) + return nil + }) + Expect(err).ToNot(HaveOccurred()) + Expect(offenders).To(BeEmpty(), + "epochs are unique and not ordered; compare them with == or !=") + }) +}) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index b4040c83e093..ab5cfa079d4c 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -151,7 +151,8 @@ func ensureEpochSequence(ctx context.Context, db *gorm.DB) error { // treat it as a version number. The sequence value on the insert path is drawn // while the tuple is built, before the row lock, so a claim that inserts after // a Release can be handed a number lower than one already issued elsewhere. -// Compare epochs for equality only; never compare them for order. +// Compare epochs for equality only; never compare them for order. A spec +// (epoch_order_test.go) fails on a < or > of an epoch in code outside the specs. // // Uniqueness is all the fence needs: Release matches owner and epoch exactly, // so a stale claim's token cannot match a live claim's row whichever way the @@ -232,6 +233,10 @@ func (r *Registry) Claim(ctx context.Context, nodeID, ownerID string) (int64, er // two statements the owner can die, and the caller would act on an owner the // second read would have rejected. The join makes the two facts one snapshot. // +// The epoch that it returns is the fence token of the claim. Compare it for +// equality with the epoch of another read or of a Claim; never for order. See +// Claim. +// // The window is InstanceLiveness rather than a parameter, which is the window // the membership loop sweeps with. A caller free to pick its own could keep // relaying to a replica the sweeper has already declared dead, or give up on @@ -295,6 +300,10 @@ func departure() map[string]any { // one that had never connected, so nothing above could tell a two-second blip // from a worker that is gone. // +// The epoch is matched for equality, and it must stay so: epochs are unique and +// not ordered (see Claim), so "the newer claim" is not a question that two +// epochs can answer. +// // Both ownerID and epoch are in the WHERE so a replica that has only just // noticed its dead socket cannot touch the claim a later reconnect established // elsewhere: it must not clear a live claim, and it must not stamp a departure From c9c6edfdd763124bfd51e0b47d3b7598c9922823 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:25:12 +0000 Subject: [PATCH 041/162] refactor(cluster): drop Registry.Live, which only specs called Nothing outside the specs read the live set through Live, and LiveInstanceIDsSQL carries the same predicate for the callers that need liveness inside another statement. The specs read the live set through that SQL now, so it stays covered. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed_carrier_test.go | 8 ++------ core/services/cluster/helpers_test.go | 11 +++++++++++ core/services/cluster/instance.go | 20 ++++---------------- core/services/cluster/instance_test.go | 12 ++++++------ core/services/cluster/membership.go | 2 +- core/services/cluster/membership_test.go | 12 ++++++------ 6 files changed, 30 insertions(+), 35 deletions(-) diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go index 76fdae9f601e..13d7c651331d 100644 --- a/core/application/distributed_carrier_test.go +++ b/core/application/distributed_carrier_test.go @@ -126,12 +126,8 @@ var _ = Describe("distributed startup and the cluster carrier", func() { It("registers every replica in the instances table and removes it on shutdown", func() { live := func() []string { GinkgoHelper() - instances, err := cluster.NewRegistry(db).Live(context.Background(), cluster.InstanceLiveness) - Expect(err).ToNot(HaveOccurred()) - ids := make([]string, 0, len(instances)) - for _, in := range instances { - ids = append(ids, in.ID) - } + var ids []string + Expect(db.Raw(cluster.LiveInstanceIDsSQL+" ORDER BY id", cluster.InstanceLiveness.Seconds()).Scan(&ids).Error).To(Succeed()) return ids } diff --git a/core/services/cluster/helpers_test.go b/core/services/cluster/helpers_test.go index 32b50f7afebc..21025250639a 100644 --- a/core/services/cluster/helpers_test.go +++ b/core/services/cluster/helpers_test.go @@ -3,6 +3,7 @@ package cluster_test import ( "context" "errors" + "time" "github.com/mudler/LocalAI/core/services/cluster" "gorm.io/gorm" @@ -22,3 +23,13 @@ func ownerRow(ctx context.Context, db *gorm.DB, nodeID string) (string, int64, e } return conn.OwnerInstanceID, conn.Epoch, nil } + +// liveInstances reads the replicas that count as alive within the window, by the +// same SQL that the claim queue joins on. +func liveInstances(ctx context.Context, db *gorm.DB, within time.Duration) ([]cluster.Instance, error) { + var out []cluster.Instance + err := db.WithContext(ctx). + Raw("SELECT * FROM instances WHERE id IN ("+cluster.LiveInstanceIDsSQL+") ORDER BY id", within.Seconds()). + Scan(&out).Error + return out, err +} diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index 2f5256179e9f..df2290aae6b3 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -44,8 +44,8 @@ func NewRegistry(db *gorm.DB) *Registry { } // Register records this replica and refreshes its LastSeen. It upserts on the -// primary key and does not delete and insert, so a concurrent Live never sees a -// live replica as missing. The readiness columns are written too, so a replica +// primary key and does not delete and insert, so a concurrent reader of the +// live set never sees a live replica as missing. The readiness columns are written too, so a replica // that registers again after it was swept does not lose what it reported. func (r *Registry) Register(ctx context.Context, id, version string, readyEpoch int64, readyReason string) error { // The database stamps last_seen and not this process. Liveness is compared @@ -113,8 +113,8 @@ func (r *Registry) ReportReady(ctx context.Context, id string, epoch int64, reas // instanceIsLive is the one predicate that decides whether a replica is alive. // Its single bind parameter is the window in seconds. Every reader of that fact -// is written with it: Live lists the rows it selects, Owner refuses an owner -// that it rejects, and ReapStale deletes its negation. Two spellings of one +// is written with it: LiveInstanceIDsSQL lists the ids it selects, Owner +// refuses an owner that it rejects, and ReapStale deletes its negation. Two spellings of one // fact drift apart, and the drift looks like a relay to a replica that one // query calls dead and another calls alive. // @@ -138,15 +138,3 @@ const instanceIsLive = `instances.last_seen > now() - make_interval(secs => ?)` // It is built from instanceIsLive and does not restate it, so the deployment // has one spelling of "alive". const LiveInstanceIDsSQL = `SELECT instances.id FROM instances WHERE ` + instanceIsLive - -// Live returns the instances whose LastSeen is newer than now minus within. -func (r *Registry) Live(ctx context.Context, within time.Duration) ([]Instance, error) { - var out []Instance - if err := r.db.WithContext(ctx). - Where(instanceIsLive, within.Seconds()). - Order("id"). - Find(&out).Error; err != nil { - return nil, fmt.Errorf("listing live instances: %w", err) - } - return out, nil -} diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go index 21944c14a8e8..45b93dc516f6 100644 --- a/core/services/cluster/instance_test.go +++ b/core/services/cluster/instance_test.go @@ -38,7 +38,7 @@ var _ = Describe("Instance registry", func() { It("registers an instance and reads it back", func() { Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) Expect(live[0].ID).To(Equal("inst-a")) @@ -57,7 +57,7 @@ var _ = Describe("Instance registry", func() { Expect(reg.Register(ctx, "inst-a", "v1", 0, "")).To(Succeed()) Expect(reg.Register(ctx, "inst-a", "v2", 0, "")).To(Succeed()) - live, err := reg.Live(ctx, time.Hour) + live, err := liveInstances(ctx, db, time.Hour) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) Expect(live[0].Version).To(Equal("v2")) @@ -69,7 +69,7 @@ var _ = Describe("Instance registry", func() { Expect(db.Model(&cluster.Instance{}).Where("id = ?", "stale"). Update("last_seen", gorm.Expr("now() - interval '10 minutes'")).Error).To(Succeed()) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(BeEmpty()) }) @@ -91,14 +91,14 @@ var _ = Describe("Instance registry", func() { Update("last_seen", gorm.Expr("now() - interval '10 minutes'")).Error).To(Succeed()) Expect(reg.Heartbeat(ctx, "revive")).To(Succeed()) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) }) It("heartbeating an unknown instance is an error and not a silent insert", func() { Expect(reg.Heartbeat(ctx, "ghost")).To(MatchError(cluster.ErrInstanceNotFound)) - live, err := reg.Live(ctx, time.Hour) + live, err := liveInstances(ctx, db, time.Hour) Expect(err).ToNot(HaveOccurred()) Expect(live).To(BeEmpty()) }) @@ -137,7 +137,7 @@ var _ = Describe("Instance registry", func() { Expect(reg.ReportReady(ctx, "inst-a", 2, "")).To(Succeed()) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) }) diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index 59e5df771ffa..7e72e901cbf7 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -415,7 +415,7 @@ func (r *Registry) Deregister(ctx context.Context, id string) error { // The protection goes one way: another replica can still reap a replica that // stalls for long enough. // -// It needs PostgreSQL, like Live, and the interval is measured on the clock of +// It needs PostgreSQL, like the reads of liveness, and the interval is measured on the clock of // the database. func (r *Registry) ReapStale(ctx context.Context, self string, within time.Duration) (instances int64, cleared int64, err error) { err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go index ebc57060c192..09684e417f29 100644 --- a/core/services/cluster/membership_test.go +++ b/core/services/cluster/membership_test.go @@ -164,7 +164,7 @@ var _ = Describe("Reaping dead replicas", func() { Expect(reg.Deregister(ctx, "leaving")).To(Succeed()) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) Expect(live[0].ID).To(Equal("staying")) @@ -280,13 +280,13 @@ var _ = Describe("Reaping dead replicas", func() { membership := cluster.NewMembership(reg, "me", "v1") Expect(membership.Start(ctx)).To(Succeed()) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) membership.Stop() - live, err = reg.Live(ctx, time.Minute) + live, err = liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(BeEmpty(), "a replica that shut down cleanly left its row behind for peers to dial") }) @@ -343,9 +343,9 @@ var _ = Describe("Reaping dead replicas", func() { Expect(membership.Start(ctx)).To(Succeed()) DeferCleanup(membership.Stop) - // Rows, not live rows: an aged-out replica drops out of Live + // Rows, not live rows: an aged-out replica drops out of the live set // immediately, and what the sweeper adds is deleting it. Asserting on - // Live here would pass with no sweeper at all. + // the live set here would pass with no sweeper at all. rows := func() int64 { var n int64 if err := db.Model(&cluster.Instance{}).Count(&n).Error; err != nil { @@ -356,7 +356,7 @@ var _ = Describe("Reaping dead replicas", func() { Expect(rows()).To(Equal(int64(2)), "the stale row is still in the table until a sweep deletes it") Eventually(rows, 3*cluster.InstanceHeartbeat, time.Second).Should(Equal(int64(1))) - live, err := reg.Live(ctx, time.Minute) + live, err := liveInstances(ctx, db, time.Minute) Expect(err).ToNot(HaveOccurred()) Expect(live).To(HaveLen(1)) Expect(live[0].ID).To(Equal("me"), "the sweeper deleted the wrong row") From 7546943ac045b4d91207097ef5fee1688e9754fc Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:25:55 +0000 Subject: [PATCH 042/162] test(tunnel): cover the claim again after a reaped replica registers A peer reaps a replica that stalled. The membership loop registers it again and asks the tunnel registry to claim the sockets that it still holds. The spec runs that whole path on a real database and fails when the reclaimer is not set. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/registry_test.go | 30 +++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/core/services/tunnel/registry_test.go b/core/services/tunnel/registry_test.go index 0e8a8224d0e5..719d6dab30f9 100644 --- a/core/services/tunnel/registry_test.go +++ b/core/services/tunnel/registry_test.go @@ -381,6 +381,36 @@ var _ = Describe("Registry", func() { Expect(id).To(Equal("replica-a")) }) + // The whole path: a peer reaps this replica, the membership loop registers + // it again, and the loop asks the registry of tunnels to claim the held + // sockets again. A replica that re-registers and does not do this serves + // a worker that the table says nobody holds. + It("claims the held tunnels again when the membership loop registers a reaped replica again", func() { + membership := cluster.NewMembership(clusterR, "replica-a", "test") + membership.SetReclaimer(registry) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + frontend, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "replica-a"). + Update("last_seen", gorm.Expr("now() - interval '1 hour'")).Error).To(Succeed()) + instances, cleared, err := clusterR.ReapStale(ctx, "replica-b", cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(instances).To(Equal(int64(1))) + Expect(cleared).To(Equal(int64(1))) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + + Eventually(func() string { + id, _ := owner("w1") + return id + }, 3*cluster.InstanceHeartbeat, 250*time.Millisecond).Should(Equal("replica-a"), + "the replica registered again and left its held tunnel unclaimed") + }) + It("lets the Detach of the attachment release the claim that Reclaim wrote", func() { frontend, _ := sessionPair(tunnel.LaneInference) token, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) From c58d6d4df44cceefbbf82d3f0994f6266af2d174 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:50:32 +0000 Subject: [PATCH 043/162] fix(worker): register again when the frontend refuses the tunnel credential Another process that registers under the same node name replaces the credential of the node, and the first worker then got 401 on every dial for ever. It logged that it needed a restart and skipped its heartbeat while the tunnel was down. The tunnel client now registers again after a 401 or a 403 and dials with the new credential. It does so once for each wait of the backoff, so a frontend that keeps refusing is not asked in a loop. The comment on the credential check claimed more than it does. The credential stops a caller from claiming a node by its ID alone. The registration token still authorises a registration under any name, as it does for the NATS credentials. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/workerregistry/credentials.go | 18 ++++ core/cli/workerregistry/credentials_test.go | 39 +++++++ core/http/endpoints/cluster/connect.go | 7 ++ core/http/endpoints/cluster/connect_test.go | 58 +++++++++++ core/services/worker/tunnel.go | 47 +++++++++ core/services/worker/tunnel_test.go | 109 ++++++++++++++++++++ core/services/worker/worker.go | 20 +++- 7 files changed, 297 insertions(+), 1 deletion(-) diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 6168876b60ef..2b052f7df4a1 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -182,6 +182,24 @@ func (m *CredentialManager) Acquire(ctx context.Context) (*RegisterResponse, err return nil, fmt.Errorf("giving up acquiring NATS credentials after %d attempts: %w", m.maxAttempts, lastReason) } +// Reregister registers once more and keeps the credentials of the answer. It +// does not wait and does not retry: the caller owns the backoff. The tunnel +// client calls it after the frontend refused its credential, which happens when +// another process registered under the same node name, or when the frontend +// lost its record of the node. A registration mints a new credential, and the +// registration token of the deployment authorises it, as it does at startup. +func (m *CredentialManager) Reregister(ctx context.Context) error { + res, err := m.register(ctx) + if err != nil { + return err + } + if res.Carrier == carrierTunnel && res.TunnelToken == "" { + return fmt.Errorf("node %s registered but the tunnel credential was not minted", res.ID) + } + m.store(res) + return nil +} + // RefreshLoop re-registers to mint a fresh JWT before the current one expires, // updating the credentials returned by Current/Provider so the NATS connection // adopts them on its next reconnect. It returns nil when ctx is cancelled or diff --git a/core/cli/workerregistry/credentials_test.go b/core/cli/workerregistry/credentials_test.go index f1f71fa5d3cf..76de4fd749eb 100644 --- a/core/cli/workerregistry/credentials_test.go +++ b/core/cli/workerregistry/credentials_test.go @@ -2,6 +2,7 @@ package workerregistry import ( "context" + "errors" "sync" "testing" "time" @@ -116,6 +117,44 @@ var _ = Describe("CredentialManager", func() { }) }) + Describe("Reregister", func() { + tunnelRes := func(token string) *RegisterResponse { + return &RegisterResponse{ID: "node-1", Status: "healthy", Carrier: "tunnel", TunnelToken: token} + } + + It("replaces the tunnel token with the one that the new registration minted", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("token-1")}, {res: tunnelRes("token-2")}}} + m := NewCredentialManager(f.fn(), false) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + + Expect(m.Reregister(context.Background())).To(Succeed()) + Expect(m.TunnelToken()).To(Equal("token-2")) + Expect(f.count()).To(Equal(2), "one registration and no retry") + }) + + It("keeps the old token and returns the error when registration fails", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("token-1")}, {err: errors.New("frontend down")}}} + m := NewCredentialManager(f.fn(), false) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + + Expect(m.Reregister(context.Background())).To(MatchError(ContainSubstring("frontend down"))) + Expect(m.TunnelToken()).To(Equal("token-1")) + Expect(f.count()).To(Equal(2)) + }) + + It("keeps the old token when the answer names the tunnel and mints none", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("token-1")}, {res: tunnelRes("")}}} + m := NewCredentialManager(f.fn(), false) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + + Expect(m.Reregister(context.Background())).ToNot(Succeed()) + Expect(m.TunnelToken()).To(Equal("token-1")) + }) + }) + Describe("Acquire on the tunnel carrier", func() { tunnelRes := func(status, token string) *RegisterResponse { return &RegisterResponse{ID: "node-1", Status: status, Carrier: "tunnel", CarrierEpoch: 4, TunnelToken: token} diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go index 2bd5231eef4f..b2bddb150ac6 100644 --- a/core/http/endpoints/cluster/connect.go +++ b/core/http/endpoints/cluster/connect.go @@ -212,6 +212,13 @@ func bearerToken(r *http.Request) (string, bool) { // node at registration, returned to that worker once and stored as this hash. // It is not TokenHash, which hashes the token that the worker registered with. // On most deployments that is the shared registration token. +// +// The credential stops a caller from claiming a node by its ID alone. It does +// not stop a holder of the registration token: that token authorises a +// registration under any node name, and a registration mints a new credential +// and replaces the old one. This is the same trust as the NATS credentials that +// registration mints. A worker whose credential was replaced registers again +// by itself. func authorizedWorker(token, storedHash string) bool { if storedHash == "" { return false diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go index cab45f4c75c6..cb0f433136c6 100644 --- a/core/http/endpoints/cluster/connect_test.go +++ b/core/http/endpoints/cluster/connect_test.go @@ -7,6 +7,8 @@ import ( "io" "net/http" "net/http/httptest" + "sync" + "sync/atomic" "time" "github.com/gorilla/websocket" @@ -17,6 +19,7 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/testutil" "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/worker" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -333,6 +336,61 @@ var _ = Describe("Connect endpoint", func() { }) }) +var _ = Describe("A worker whose credential was rotated on the frontend", func() { + It("registers again and gets its tunnel back without a restart", func() { + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR := cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + tunnels := tunnel.NewRegistry(clusterR, "replica-a") + + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf("registration")} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + var mu sync.Mutex + current := "credential-1" + Expect(nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(current))).To(Succeed()) + held := func() string { mu.Lock(); defer mu.Unlock(); return current } + + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(nodeReg, tunnels)) + srv := httptest.NewServer(e) + DeferCleanup(srv.Close) + + // Another process registers under the same name and replaces the + // credential on the frontend. This worker still holds credential-1. + mu.Lock() + current = "credential-2-of-another-process" + Expect(nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(current))).To(Succeed()) + mu.Unlock() + workerToken := "credential-1" + var reregistered atomic.Int32 + w, err := worker.StartTunnel(ctx, worker.TunnelConfig{ + FrontendURL: srv.URL, + NodeID: node.ID, + Token: func() string { mu.Lock(); defer mu.Unlock(); return workerToken }, + Services: map[string]worker.LocalService{}, + Reauthorize: func(ctx context.Context) error { + // What registration does on the frontend: mint, store the hash, + // return the value once. + mu.Lock() + defer mu.Unlock() + reregistered.Add(1) + current = "credential-3" + workerToken = current + return nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(current)) + }, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(w.Close) + Eventually(reregistered.Load, 20*time.Second).Should(BeNumerically(">=", 1)) + Eventually(w.Connected, 20*time.Second).Should(BeTrue()) + Expect(held()).To(Equal("credential-3")) + }) +}) + var _ = Describe("The path of the connect route", func() { It("is the one that the global auth middleware lets through", func() { Expect(auth.ClusterConnectPath).To(Equal(tunnel.ConnectPath)) diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index 99652e1a3dd1..0bb8bf4d03ea 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -106,6 +106,15 @@ type TunnelConfig struct { // entry is refused. See accept. Services map[string]LocalService + // Reauthorize registers the node again, so that Token returns a credential + // that the frontend accepts. It is called after a dial that the frontend + // refused with 401 or 403, and at most once for each wait of the backoff. + // Another process that registers under the same node name, or a frontend that + // lost its record of the node, makes the credential of this worker useless, + // and only a registration mints a new one. A nil Reauthorize leaves the + // worker to retry the credential it has. + Reauthorize func(ctx context.Context) error + // The fields below are for the specs. They are not exported, so that a caller // cannot reach them. sleep func(ctx context.Context, d time.Duration) error @@ -126,6 +135,8 @@ type Tunnel struct { services map[string]LocalService dialer *websocket.Dialer + reauthorize func(ctx context.Context) error + headerTimeout time.Duration noBulk bool sleep func(ctx context.Context, d time.Duration) error @@ -224,6 +235,7 @@ func StartTunnel(ctx context.Context, cfg TunnelConfig) (*Tunnel, error) { token: cfg.Token, services: services, dialer: dialer, + reauthorize: cfg.Reauthorize, headerTimeout: cmp.Or(cfg.headerTimeout, tunnelHeaderTimeout), noBulk: cfg.noBulk, sleep: cfg.sleep, @@ -279,12 +291,47 @@ func (t *Tunnel) run(ctx context.Context) { delay := tunnelBackoffDelay(attempt) t.logSessionEnded(err, attempt, delay) + t.reauthorizeAfter(ctx, err) if err := t.sleep(ctx, delay); err != nil { return } } } +// reauthorizeAfter registers the node again when the frontend refused the +// credential of the dial. +// +// It runs once for each pass of the loop, and the loop waits with a growing +// backoff after it, so a frontend that keeps refusing is asked no more often +// than the backoff allows. Two workers that share a node name rotate each +// other's credential on every registration; the backoff bounds that contest to +// one registration for each wait and does not end it. The operator fixes the +// duplicate name. +func (t *Tunnel) reauthorizeAfter(ctx context.Context, err error) { + if t.reauthorize == nil || !refusedCredential(err) { + return + } + if rerr := t.reauthorize(ctx); rerr != nil { + if ctx.Err() == nil { + xlog.Warn("Registering again after the frontend refused the tunnel credential failed", "node", t.nodeID, "error", rerr) + } + return + } + xlog.Info("Registered again after the frontend refused the tunnel credential", "node", t.nodeID) +} + +// refusedCredential reports whether the frontend answered a dial with 401 or +// 403. Both can be cured by a new registration: 401 is a credential that the +// frontend does not know, and 403 is a node whose record the registration +// updates. +func refusedCredential(err error) bool { + var dialErr *tunnelDialError + if !errors.As(err, &dialErr) { + return false + } + return dialErr.status == http.StatusUnauthorized || dialErr.status == http.StatusForbidden +} + // connectAndServe dials the inference lane, serves streams until the session // ends, and leaves nothing running behind it. func (t *Tunnel) connectAndServe(ctx context.Context) error { diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go index 3aa10e8c14e7..c86533a53cc3 100644 --- a/core/services/worker/tunnel_test.go +++ b/core/services/worker/tunnel_test.go @@ -59,6 +59,12 @@ type fakeFrontend struct { // bulkStatus, when it is not zero, is the status that a bulk dial gets // instead of an upgrade. bulkStatus atomic.Int32 + // requiredToken, when it is set, is the only credential that a dial may + // present. Any other gets 401, as the connect endpoint answers a node whose + // credential was rotated. + requiredToken atomic.Pointer[string] + // denyStatus, when it is not zero, is the status that every dial gets. + denyStatus atomic.Int32 } func newFakeFrontend(closeAtOnce bool) *fakeFrontend { @@ -87,6 +93,14 @@ func newFakeFrontend(closeAtOnce bool) *fakeFrontend { }: default: } + if want := f.requiredToken.Load(); want != nil && strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") != *want { + w.WriteHeader(http.StatusUnauthorized) + return + } + if status := int(f.denyStatus.Load()); status != 0 { + w.WriteHeader(status) + return + } if status := int(f.bulkStatus.Load()); lane == tunnel.LaneBulk && status != 0 { w.WriteHeader(status) return @@ -554,6 +568,101 @@ var _ = Describe("Worker tunnel client", func() { Expect(first.nodeID).To(Equal("node-1")) }) + Describe("after the frontend refused the credential", func() { + // The frontend rotates the credential of the node, as it does when a + // second process registers under the same name. The worker holds the + // old value and must get a new one without a restart. + It("registers again and connects with the new credential", func() { + frontend = newFakeFrontend(false) + current := "rotated-by-another-process" + frontend.requiredToken.Store(¤t) + var token atomic.Pointer[string] + stale := "stale" + token.Store(&stale) + var registrations atomic.Int32 + start(func(c *TunnelConfig) { + c.noBulk = true + c.Token = func() string { return *token.Load() } + c.Reauthorize = func(context.Context) error { + registrations.Add(1) + token.Store(¤t) + return nil + } + c.sleep = func(ctx context.Context, _ time.Duration) error { return ctx.Err() } + }) + Eventually(frontend.inference, "10s").Should(Receive()) + Expect(registrations.Load()).To(Equal(int32(1))) + Eventually(tun.Connected, "10s").Should(BeTrue()) + }) + + It("registers again once for each wait and never in a tight loop", func() { + frontend = newFakeFrontend(false) + frontend.denyStatus.Store(http.StatusForbidden) + var registrations atomic.Int32 + // Unbuffered: the loop cannot start a pass until the spec has read the + // wait of the last one. + waits := make(chan time.Duration) + start(func(c *TunnelConfig) { + c.noBulk = true + c.Reauthorize = func(context.Context) error { registrations.Add(1); return nil } + c.sleep = func(ctx context.Context, d time.Duration) error { + select { + case waits <- d: + return nil + case <-ctx.Done(): + return ctx.Err() + } + } + }) + var d time.Duration + for range 5 { + Eventually(waits, "10s").Should(Receive(&d)) + Expect(d).To(BeNumerically(">", 0)) + } + // One registration for each pass, and every pass waits. + Expect(int(registrations.Load())).To(BeNumerically("<=", 6)) + Expect(int(registrations.Load())).To(BeNumerically(">=", 5)) + }) + + It("keeps trying with the backoff when registering again fails", func() { + frontend = newFakeFrontend(false) + frontend.denyStatus.Store(http.StatusUnauthorized) + var registrations atomic.Int32 + start(func(c *TunnelConfig) { + c.noBulk = true + c.Reauthorize = func(context.Context) error { + registrations.Add(1) + return errors.New("frontend down") + } + c.sleep = func(ctx context.Context, _ time.Duration) error { return ctx.Err() } + }) + Eventually(registrations.Load, "10s").Should(BeNumerically(">=", 3)) + Expect(tun.Connected()).To(BeFalse()) + }) + + It("does not register again for a refusal that a registration does not cure", func() { + frontend = newFakeFrontend(false) + frontend.denyStatus.Store(http.StatusServiceUnavailable) + var registrations atomic.Int32 + dials := make(chan struct{}, 64) + start(func(c *TunnelConfig) { + c.noBulk = true + c.Reauthorize = func(context.Context) error { registrations.Add(1); return nil } + c.sleep = func(ctx context.Context, _ time.Duration) error { + select { + case dials <- struct{}{}: + default: + } + return ctx.Err() + } + }) + for range 3 { + Eventually(dials, "10s").Should(Receive()) + } + Expect(registrations.Load()).To(BeZero()) + }) + }) + It("does not dial with no credential, and says so", func() { frontend = newFakeFrontend(false) start(func(c *TunnelConfig) { diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 8bf0401e96f1..c5cd31c517e5 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -10,6 +10,7 @@ import ( "os" "os/signal" "path/filepath" + "sync/atomic" "syscall" "time" @@ -125,8 +126,22 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // registration. With static NATS credentials there is no manager, and a // frontend that answers with the tunnel gives the token once. tunnelToken := credMgr.TunnelToken + reauthorize := credMgr.Reregister if staticNATS { - tunnelToken = func() string { return res.TunnelToken } + var staticToken atomic.Pointer[string] + staticToken.Store(&res.TunnelToken) + tunnelToken = func() string { return *staticToken.Load() } + reauthorize = func(ctx context.Context) error { + again, err := regClient.RegisterFull(ctx, registrationBody) + if err != nil { + return err + } + if again.TunnelToken == "" { + return fmt.Errorf("node %s registered but the tunnel credential was not minted", again.ID) + } + staticToken.Store(&again.TunnelToken) + return nil + } } xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cfg.RegisterTo, "carrier", cmp.Or(res.Carrier, "nats")) @@ -187,6 +202,9 @@ func Run(ctx *cliContext.Context, cfg *Config) error { FrontendURL: cfg.RegisterTo, NodeID: nodeID, Token: tunnelToken, + // A frontend that refuses the credential is told who this worker is + // again, so that a rotated credential does not leave it unreachable. + Reauthorize: reauthorize, // Built by tunnelServices and not inline, so that the routing table, // which is the security boundary of the tunnel, can be reached from a // spec without starting a worker. From 99523406dd2f7da1cee05625049c37273d0c2c76 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 7 Oct 2026 23:52:23 +0000 Subject: [PATCH 044/162] fix(worker): read only an unreachable target as a reason to reap A failed local dial answered "unavailable" for every error that was not a timeout. The frontend counts that answer as evidence about a backend and can remove the rows of a loaded model. A worker that ran out of file descriptors, buffers, memory or local ports would then report its own shortage as a missing backend. The classifier now keeps an allow list: connection refused, host or network unreachable, and connection reset or abort mean that the target did not answer. Every other error means that the worker learned nothing. Windows lists the WSA values as well. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/errno_unix.go | 15 +++++++ core/services/worker/errno_windows.go | 21 ++++++++++ core/services/worker/tunnel.go | 59 +++++++++++++++------------ core/services/worker/tunnel_test.go | 40 +++++++++++++++++- 4 files changed, 109 insertions(+), 26 deletions(-) create mode 100644 core/services/worker/errno_unix.go create mode 100644 core/services/worker/errno_windows.go diff --git a/core/services/worker/errno_unix.go b/core/services/worker/errno_unix.go new file mode 100644 index 000000000000..b6512fe384a1 --- /dev/null +++ b/core/services/worker/errno_unix.go @@ -0,0 +1,15 @@ +//go:build !windows + +package worker + +import "syscall" + +// unreachableErrnos are the errors of a dial that say the target did not +// answer. See classifyServiceFailure. +var unreachableErrnos = []error{ + syscall.ECONNREFUSED, + syscall.EHOSTUNREACH, + syscall.ENETUNREACH, + syscall.ECONNRESET, + syscall.ECONNABORTED, +} diff --git a/core/services/worker/errno_windows.go b/core/services/worker/errno_windows.go new file mode 100644 index 000000000000..5fe24797ee47 --- /dev/null +++ b/core/services/worker/errno_windows.go @@ -0,0 +1,21 @@ +//go:build windows + +package worker + +import "syscall" + +// unreachableErrnos are the errors of a dial that say the target did not +// answer. See classifyServiceFailure. Windows reports the WSA values and not +// the invented values of package syscall, so both are listed. +var unreachableErrnos = []error{ + syscall.ECONNREFUSED, + syscall.EHOSTUNREACH, + syscall.ENETUNREACH, + syscall.ECONNRESET, + syscall.ECONNABORTED, + syscall.Errno(10061), // WSAECONNREFUSED + syscall.Errno(10065), // WSAEHOSTUNREACH + syscall.Errno(10051), // WSAENETUNREACH + syscall.Errno(10054), // WSAECONNRESET + syscall.Errno(10053), // WSAECONNABORTED +} diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index 0bb8bf4d03ea..ae7e77455dde 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -563,31 +563,25 @@ func (t *Tunnel) accept(ctx context.Context, stream net.Conn) (net.Conn, bool) { // cannot work. The second reported as the first makes it give up on a backend // that is starting. // -// The default is "unavailable", and that is on purpose, because the frontend -// acts on that code. A dial to the process that failed is the nearest thing to -// evidence that the worker can produce. The direction of a wrong classification -// decides which mistake is made. A wrong eviction is active: it deletes rows -// and shuts down a model that is loaded. A wrong retention is passive, is -// limited to one slot of one replica, and clears at a restart or an eviction. -// Because of this, the exceptions below are a deny list of causes that are not -// the target answering, and not an allow list of causes that are. A deny list -// that misses a cause fails loudly. An allow list that misses one fails -// silently, for ever. +// The default is "this worker learned nothing", and "unavailable" is an allow +// list. The frontend acts on "unavailable": it counts as evidence that can remove +// the rows of a model. A wrong removal is active. It deletes rows and shuts down +// a model that is loaded. A wrong retention is passive, is limited to one slot of +// one replica, and clears at a restart or an eviction. Because of this, only an +// error that says the target did not answer counts. A dial can also fail for a +// reason in this process or on this host: no file descriptors (EMFILE, ENFILE), +// no buffers or memory (ENOBUFS, ENOMEM), no local port (EADDRNOTAVAIL), or a +// local rule (EACCES, EPERM). The target was never asked in those cases, and a +// reading of them as "the backend is gone" would remove a healthy model when +// this worker is short of descriptors. An error that nobody classified is the +// same: it can be silent for ever, and that is the cheaper mistake. // -// The causes that are exempt: -// -// - The end of the context. It is the context of the session, and it ends -// while stream goroutines are running, so this worker is closing its -// tunnel and will connect again. -// - The own deadline of this process (os.ErrDeadlineExceeded). -// - Anything else that says of itself that it is a timeout. On a dial this -// includes ETIMEDOUT and EAGAIN. EAGAIN is this worker running out of a -// resource, and that is not about the target. ETIMEDOUT from connect(2) is -// an observation of the target, but it is "the handshake did not finish", -// which is a listener that is blocked or full and not one that is absent. -// An eviction of a backend that is overloaded is what this design must -// prevent. ECONNREFUSED, the shape of a process that is gone, is not a -// timeout and still gives "unavailable". +// The errors that mean the target did not answer are in targetUnreachable. The +// end of the context, the own deadline of this process and anything that says of +// itself that it is a timeout are never among them. The context ends while +// stream goroutines run, because this worker is closing its tunnel. A timeout of +// connect(2) means that the handshake did not finish, which is a listener that is +// blocked or full and not one that is absent. // // It must never give the unknown-tag refusal: a tag that the worker serves does // not stop being served because one dial failed. @@ -598,7 +592,22 @@ func classifyServiceFailure(err error) error { if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || reportsTimeout(err) { return fmt.Errorf("%w: %v", tunnel.ErrStreamNotServed, err) } - return fmt.Errorf("%w: %v", tunnel.ErrStreamTargetUnavailable, err) + if targetUnreachable(err) { + return fmt.Errorf("%w: %v", tunnel.ErrStreamTargetUnavailable, err) + } + return fmt.Errorf("%w: %v", tunnel.ErrStreamNotServed, err) +} + +// targetUnreachable reports whether err says that the dial reached the network +// stack and the target did not answer: nothing listens, or the route is gone, or +// the peer reset the connection. +func targetUnreachable(err error) bool { + for _, errno := range unreachableErrnos { + if errors.Is(err, errno) { + return true + } + } + return false } // reportsTimeout reports whether err says of itself that it is a timeout. diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go index c86533a53cc3..1ce771c56e2a 100644 --- a/core/services/worker/tunnel_test.go +++ b/core/services/worker/tunnel_test.go @@ -10,9 +10,11 @@ import ( "net" "net/http" "net/http/httptest" + "os" "strconv" "strings" "sync/atomic" + "syscall" "time" "github.com/mudler/LocalAI/core/services/nodes" @@ -335,7 +337,7 @@ var _ = Describe("Worker tunnel client", func() { frontend = newFakeFrontend(false) start(func(c *TunnelConfig) { c.Services[tunnel.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { - return nil, errors.New("connection refused") + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: os.NewSyscallError("connect", syscall.ECONNREFUSED)} } }) _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:1") @@ -414,6 +416,42 @@ var _ = Describe("Worker tunnel client", func() { Expect(err).To(MatchError(tunnel.ErrStreamNotServed)) }) + DescribeTable("classifies the errno of a failed local dial", + func(errno syscall.Errno, wantEvidence bool) { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[tunnel.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { + return nil, &net.OpError{Op: "dial", Net: "tcp", Err: os.NewSyscallError("connect", errno)} + } + }) + _, err := open(session(), tunnel.StreamTagGRPC, "127.0.0.1:41000") + if wantEvidence { + Expect(err).To(MatchError(tunnel.ErrStreamTargetUnavailable)) + } else { + Expect(err).To(MatchError(tunnel.ErrStreamNotServed)) + } + Expect(tunnel.IsWorkerAnswer(err)).To(Equal(wantEvidence)) + }, + // The target did not answer: the backend is gone or not there. + Entry("connection refused", syscall.ECONNREFUSED, true), + Entry("host unreachable", syscall.EHOSTUNREACH, true), + Entry("network unreachable", syscall.ENETUNREACH, true), + Entry("connection reset", syscall.ECONNRESET, true), + // This process or this host ran short of something. The target was + // never asked, so nothing is learned about it. + Entry("too many open files in this process", syscall.EMFILE, false), + Entry("too many open files on this host", syscall.ENFILE, false), + Entry("no buffer space", syscall.ENOBUFS, false), + Entry("out of memory", syscall.ENOMEM, false), + Entry("no local port or address left", syscall.EADDRNOTAVAIL, false), + Entry("permission denied by a local rule", syscall.EACCES, false), + Entry("operation not permitted", syscall.EPERM, false), + ) + + It("learns nothing from an error that it does not recognise", func() { + Expect(classifyServiceFailure(errors.New("something new went wrong"))).To(MatchError(tunnel.ErrStreamNotServed)) + }) + DescribeTable("keeps a classification that the local service already made", func(classified error, wantEvidence bool) { frontend = newFakeFrontend(false) From 1915690bccf0b4e532b9fb40ba9017b008ba3c3d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:00:05 +0000 Subject: [PATCH 045/162] fix(tunnel): reset the streams that a worker opens on the frontend session The frontend never reads a stream that the worker opens, and yamux accepted up to 1000 of them. Each held its window of unread data for as long as the session lived: about 64 MiB for the inference lane and about 1 GiB for the bulk lane. A worker that misbehaves could make the frontend hold that memory. The server side of a session now allows no incoming stream, so yamux resets each one. The worker side keeps the default because it accepts the streams of the frontend. The latency and throughput benchmarks sent their transfer on a stream that the worker opened. They now open the stream on the frontend and send the data from the worker, which is what a download does. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/bench_test.go | 38 ++++++++++------- core/services/tunnel/latency_test.go | 61 ++++++++++++++++++---------- core/services/tunnel/session.go | 12 +++++- core/services/tunnel/session_test.go | 51 +++++++++++++++++++++++ 4 files changed, 126 insertions(+), 36 deletions(-) diff --git a/core/services/tunnel/bench_test.go b/core/services/tunnel/bench_test.go index ec3e40d697c2..efc2d9baf577 100644 --- a/core/services/tunnel/bench_test.go +++ b/core/services/tunnel/bench_test.go @@ -3,6 +3,7 @@ package tunnel_test import ( "context" "io" + "net" "net/http" "net/http/httptest" "testing" @@ -64,35 +65,44 @@ func BenchmarkTransfer(b *testing.B) { for _, dir := range []string{"push", "pull"} { b.Run(dir+"/"+buf.name, func(b *testing.B) { frontend, worker := benchPair(b, buf.up, buf.dial, tunnel.LaneInference) - sender, receiver := worker, frontend + // The frontend opens every stream, as it does in production, and + // the direction is the direction of the data. "push" is the + // worker sending to the frontend. + chunk := make([]byte, 256<<10) + write := func(st net.Conn) error { + for sent := 0; sent < size; sent += len(chunk) { + if _, err := st.Write(chunk); err != nil { + return err + } + } + return st.Close() + } + read := func(st net.Conn) error { + _, err := io.Copy(io.Discard, st) + _ = st.Close() + return err + } + opener, acceptor := read, write if dir == "pull" { - sender, receiver = frontend, worker + opener, acceptor = write, read } - chunk := make([]byte, 256<<10) b.SetBytes(size) b.ResetTimer() for range b.N { done := make(chan error, 1) go func() { - st, err := receiver.AcceptStream() + st, err := worker.AcceptStream() if err != nil { done <- err return } - _, err = io.Copy(io.Discard, st) - _ = st.Close() - done <- err + done <- acceptor(st) }() - st, err := sender.OpenStream(context.Background()) + st, err := frontend.OpenStream(context.Background()) if err != nil { b.Fatal(err) } - for sent := 0; sent < size; sent += len(chunk) { - if _, err := st.Write(chunk); err != nil { - b.Fatal(err) - } - } - if err := st.Close(); err != nil { + if err := opener(st); err != nil { b.Fatal(err) } if err := <-done; err != nil { diff --git a/core/services/tunnel/latency_test.go b/core/services/tunnel/latency_test.go index 6dd4be37dea7..9c5c85c5414a 100644 --- a/core/services/tunnel/latency_test.go +++ b/core/services/tunnel/latency_test.go @@ -92,6 +92,41 @@ func tcpService(handle func(net.Conn)) string { return lis.Addr().String() } +// sourceOrEchoTCP is the worker side of the transfer specs. The frontend opens +// every stream, as it does in production. A stream that starts with 's' asks for +// transferSize bytes in the direction from the worker to the frontend, and any +// other stream is a probe that is echoed. +func sourceOrEchoTCP(c net.Conn) { + defer func() { _ = c.Close() }() + kind := make([]byte, 1) + if _, err := io.ReadFull(c, kind); err != nil { + return + } + if kind[0] == 's' { + _ = pushBytes(c, transferSize) + return + } + if _, err := c.Write(kind); err != nil { + return + } + _, _ = io.Copy(c, c) +} + +// pullTransfer opens a stream from the frontend, asks the worker for the +// transfer and reads it to the end. +func pullTransfer(frontend *tunnel.Session) error { + st, err := frontend.OpenStream(context.Background()) + if err != nil { + return err + } + defer func() { _ = st.Close() }() + if _, err := st.Write([]byte{'s'}); err != nil { + return err + } + _, err = io.Copy(io.Discard, st) + return err +} + func pushBytes(w io.Writer, size int) error { chunk := make([]byte, 256<<10) for sent := 0; sent < size; sent += len(chunk) { @@ -216,7 +251,6 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma bulkName = "bulk lane: transfer on the bulk session, probe on the inference session" ) results := map[string]probeStats{} - // measure runs the transfer while the probes run, and records the stats of // the probes. measure := func(name string, transfer func() error, probe func() (time.Duration, error)) { @@ -285,18 +319,10 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma var frontend *tunnel.Session Eventually(f.inference).Should(Receive(&frontend)) DeferCleanup(func() { _ = frontend.Close() }) - serveStreams(worker, echoTCP) - serveStreams(frontend, sinkTCP) + serveStreams(worker, sourceOrEchoTCP) measure(sharedName, - func() error { - st, err := worker.OpenStream(context.Background()) - if err != nil { - return err - } - defer func() { _ = st.Close() }() - return pushBytes(st, transferSize) - }, + func() error { return pullTransfer(frontend) }, streamProbe(frontend)) }) @@ -311,18 +337,11 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma Eventually(f.inference).Should(Receive(&inferenceFrontend)) Eventually(f.bulk).Should(Receive(&bulkFrontend)) DeferCleanup(func() { _ = inferenceFrontend.Close(); _ = bulkFrontend.Close() }) - serveStreams(inferenceWorker, echoTCP) - serveStreams(bulkFrontend, sinkTCP) + serveStreams(inferenceWorker, sourceOrEchoTCP) + serveStreams(bulkWorker, sourceOrEchoTCP) measure(bulkName, - func() error { - st, err := bulkWorker.OpenStream(context.Background()) - if err != nil { - return err - } - defer func() { _ = st.Close() }() - return pushBytes(st, transferSize) - }, + func() error { return pullTransfer(bulkFrontend) }, streamProbe(inferenceFrontend)) }) diff --git a/core/services/tunnel/session.go b/core/services/tunnel/session.go index 91772ae87967..9cc9f8e2ff0f 100644 --- a/core/services/tunnel/session.go +++ b/core/services/tunnel/session.go @@ -81,8 +81,18 @@ type Session struct { // ServerSession starts the frontend side of a session on ws. The frontend owns // the streams with even numbers and opens the streams. It does not accept any, // because only the frontend asks and the worker answers. +// +// The session allows no incoming stream, so yamux resets every stream that the +// worker opens. Nothing reads such a stream. A session that accepted it would +// keep its window of unread data for as long as the session lives, 256 KiB for +// each stream on the inference lane and up to 4 MiB on the bulk lane, until the +// default limit of 1000 streams: about 64 MiB and 1 GiB for a worker that +// misbehaves. The worker side keeps the default, because it accepts the +// streams of the frontend. func ServerSession(ws *websocket.Conn, lane Lane) (*Session, error) { - s, err := yamux.Server(websocketConn(ws), sessionConfig(lane), nil) + cfg := sessionConfig(lane) + cfg.MaxIncomingStreams = 0 + s, err := yamux.Server(websocketConn(ws), cfg, nil) if err != nil { return nil, fmt.Errorf("starting the server side of a tunnel session: %w", err) } diff --git a/core/services/tunnel/session_test.go b/core/services/tunnel/session_test.go index 509de8e61728..e579e89739e9 100644 --- a/core/services/tunnel/session_test.go +++ b/core/services/tunnel/session_test.go @@ -2,7 +2,9 @@ package tunnel_test import ( "context" + "errors" "io" + "net" "os" "path/filepath" "slices" @@ -53,6 +55,55 @@ var _ = Describe("Buffers of the websocket", func() { }) }) +var _ = Describe("Streams that the worker opens", func() { + // The frontend reads no stream that the worker opens. A stream that the + // session accepted would hold its window of unread data for as long as the + // session lives, and a worker that opens many of them would make the + // frontend hold memory for each. + for _, lane := range []tunnel.Lane{tunnel.LaneInference, tunnel.LaneBulk} { + It("resets every one of them on the "+string(lane)+" lane and accepts none", func() { + server, client := sessionPair(lane) + accepted := make(chan struct{}, 1) + go func() { + if _, err := server.AcceptStream(); err == nil { + accepted <- struct{}{} + } + }() + + // More streams than the default limit of yamux for incoming streams. + const streams = 1100 + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + opened := make([]net.Conn, 0, streams) + for range streams { + st, err := client.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("unread")) + Expect(err).ToNot(HaveOccurred()) + opened = append(opened, st) + } + for i, st := range opened { + Expect(st.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err := st.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred(), "stream %d", i) + var netErr net.Error + if errors.As(err, &netErr) { + Expect(netErr.Timeout()).To(BeFalse(), "stream %d was not reset: %v", i, err) + } + } + Consistently(accepted, "200ms").ShouldNot(Receive()) + + // The frontend still opens streams to the worker. + st, err := server.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + got, err := client.AcceptStream() + Expect(err).ToNot(HaveOccurred()) + Expect(got.Close()).To(Succeed()) + }) + } +}) + var _ = Describe("Lanes", func() { It("reads the lane parameter, and reads an empty one as the inference lane", func() { lane, err := tunnel.ParseLane("") From e63f86a4e16c9b2ab647a31fc65388cce99dc7e9 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:00:05 +0000 Subject: [PATCH 046/162] test(tunnel): run the latency benchmark only when its label is asked for The specs compare the delay of a probe during a transfer, and a loaded runner moves the numbers by more than the margin. They ran in the default suite and could fail an unrelated change. They now skip unless the label filter names benchmark. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/latency_test.go | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/core/services/tunnel/latency_test.go b/core/services/tunnel/latency_test.go index 9c5c85c5414a..2043b7730eb9 100644 --- a/core/services/tunnel/latency_test.go +++ b/core/services/tunnel/latency_test.go @@ -7,6 +7,7 @@ import ( "net/http" "net/http/httptest" "slices" + "strings" "time" "github.com/mudler/LocalAI/core/services/tunnel" @@ -251,6 +252,16 @@ var _ = Describe("Delay of a small call during a large transfer", Label("benchma bulkName = "bulk lane: transfer on the bulk session, probe on the inference session" ) results := map[string]probeStats{} + + // The specs compare latencies, and a loaded runner moves them by more than + // the margin. They run when the label is asked for: + // go run github.com/onsi/ginkgo/v2/ginkgo --label-filter=benchmark ./core/services/tunnel + BeforeAll(func() { + if !strings.Contains(GinkgoLabelFilter(), "benchmark") { + Skip("a latency benchmark; run it with --label-filter=benchmark") + } + }) + // measure runs the transfer while the probes run, and records the stats of // the probes. measure := func(name string, transfer func() error, probe func() (time.Duration, error)) { From f38b266f490ee65ab0be0042a03a78343ec9de84 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:07:17 +0000 Subject: [PATCH 047/162] feat(tunnel): end held tunnels on demand and let a bulk open refuse the fallback Registry.Disconnect ends the tunnel that a replica holds for one node and releases its claim. Registry.Close does that for every node and refuses later attaches, so a replica that shuts down does not leave claims that name it. Open takes an option, WithoutFallback. A caller that moves a large file asks for the bulk lane and gets ErrNoBulkSession when the node has none, where it used to share the inference lane silently and delay every small call on it. The default stays the fallback. Held had no caller outside the specs. The specs ask Holds. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/registry.go | 137 +++++++++++++++++++++----- core/services/tunnel/registry_test.go | 130 ++++++++++++++++++++++-- 2 files changed, 238 insertions(+), 29 deletions(-) diff --git a/core/services/tunnel/registry.go b/core/services/tunnel/registry.go index 6450cbbae323..ad87b0a43b9c 100644 --- a/core/services/tunnel/registry.go +++ b/core/services/tunnel/registry.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "net" - "sort" "sync" "time" @@ -27,6 +26,17 @@ const ConnectPath = "/api/cluster/connect" // worker that is right here, and absence must never stand in for a failure. var ErrNotOwner = errors.New("tunnel: this replica does not hold the tunnel for that node") +// ErrRegistryClosed means that the registry was closed and takes no new tunnel. +// A replica that shuts down closes it, and a worker that dials during the +// shutdown is refused and connects to another replica. +var ErrRegistryClosed = errors.New("tunnel: the registry is closed") + +// ErrNoBulkSession means that this replica holds the tunnel of the node, and +// the node holds no bulk session on it now. It is returned only to a caller that +// asked for no fallback. It is a fact about the lane and not about the node: the +// worker dials the bulk lane again within moments, and the caller may try again. +var ErrNoBulkSession = errors.New("tunnel: the node holds no bulk session on this replica") + // releaseTimeout bounds the release in Detach. Detach runs in the goroutine that // has just seen a session end, and that goroutine must not wait for a database // that went away with the session. @@ -54,6 +64,7 @@ type Registry struct { // bulkTokens numbers the attachments of the bulk lane. A bulk lane has no // claim of its own, so its attachment needs its own identity for Detach. bulkTokens int64 + closed bool } // heldTunnel is the tunnel of one node. @@ -156,6 +167,12 @@ func (t *Registry) Attach(ctx context.Context, nodeID string, lane Lane, sess *S if sess == nil { return 0, fmt.Errorf("attaching tunnel for node %q: no session", nodeID) } + t.mu.Lock() + closed := t.closed + t.mu.Unlock() + if closed { + return 0, fmt.Errorf("attaching tunnel for node %q: %w", nodeID, ErrRegistryClosed) + } switch lane { case LaneInference: return t.attachInference(ctx, nodeID, sess) @@ -185,6 +202,13 @@ func (t *Registry) attachInference(ctx context.Context, nodeID string, sess *Ses } t.mu.Lock() + if t.closed { + // Close ran during the claim. The row is released again and the + // session is not stored. + t.mu.Unlock() + t.releaseClaim(nodeID, epoch) + return 0, ErrRegistryClosed + } previous = t.tunnels[nodeID] t.tunnels[nodeID] = &heldTunnel{inference: sess, token: epoch, claim: epoch} t.mu.Unlock() @@ -278,10 +302,15 @@ func (t *Registry) Detach(nodeID string, lane Lane, token int64) { // belong to a request or a process that set the tunnel up, and either can be // cancelled by now, which would leave the row behind on every ordinary // disconnect. + t.releaseClaim(nodeID, claim) +} + +// releaseClaim releases the row of a claim. The claim, not the token: the row +// carries the last epoch that was claimed for an attachment, and the release must +// match the row exactly. +func (t *Registry) releaseClaim(nodeID string, claim int64) { ctx, cancel := context.WithTimeout(context.Background(), releaseTimeout) defer cancel() - // The claim, not the token. The row carries the last epoch that was claimed - // for this attachment, and the release must match the row exactly. if err := t.reg.Release(ctx, nodeID, t.selfID, claim); err != nil { if errors.Is(err, cluster.ErrNoConnection) { xlog.Debug("worker tunnel claim was already replaced", "node", nodeID, "epoch", claim) @@ -292,12 +321,76 @@ func (t *Registry) Detach(nodeID string, lane Lane, token int64) { } } +// Disconnect ends the tunnel that this replica holds for a node, and reports +// whether it held one. +// +// It is for a node that must lose its tunnel now: it was deleted or deregistered, +// or its credential was replaced. The credential is checked when a worker dials +// and never again, so a session that was open when the credential changed would +// otherwise go on carrying traffic for a node that has no right to it. The +// worker dials again, and a node that is gone or holds another credential is +// refused at the dial. +// +// The claim is released here and not left to the goroutine that accepted the +// session, so that a caller that asks for the owner of the node right after sees +// no owner. That goroutine detaches when its session ends and finds the entry +// gone. +func (t *Registry) Disconnect(nodeID string) bool { + t.mu.Lock() + held, ok := t.tunnels[nodeID] + if !ok { + t.mu.Unlock() + return false + } + token, inference := held.token, held.inference + t.mu.Unlock() + + t.Detach(nodeID, LaneInference, token) + _ = inference.Close() + return true +} + +// Close ends every tunnel that this replica holds and releases every claim. It +// refuses new tunnels after it. A replica calls it when it shuts down, before it +// leaves the instances table, so that a peer that asks who owns a worker does +// not name a replica that is gone. It can be called twice. +func (t *Registry) Close() { + t.mu.Lock() + t.closed = true + nodes := make([]string, 0, len(t.tunnels)) + for nodeID := range t.tunnels { + nodes = append(nodes, nodeID) + } + t.mu.Unlock() + + for _, nodeID := range nodes { + t.Disconnect(nodeID) + } +} + +// OpenOption changes how Open picks a session. +type OpenOption func(*openOptions) + +type openOptions struct { + noFallback bool +} + +// WithoutFallback makes a request for the bulk lane fail with ErrNoBulkSession +// when the node has no bulk session, and not use the inference lane. A caller +// that moves a large file asks for it: a transfer on the inference lane delays +// every small call on it, which is what the bulk lane exists to prevent. It has +// no effect on a request for the inference lane. +func WithoutFallback() OpenOption { + return func(o *openOptions) { o.noFallback = true } +} + // Open returns a stream to the worker over the tunnel that this replica holds. // // The bulk lane falls back to the inference lane when the node has no bulk -// session. A worker that predates the bulk lane never has one, and a worker -// that is dialling it again has none for a moment. A transfer that works on the -// session that is shared is better than a transfer that fails. +// session, unless the caller passes WithoutFallback. A worker that predates the +// bulk lane never has one, and a worker that is dialling it again has none for a +// moment. For a small request, a call that works on the session that is shared is +// better than a call that fails. // // ErrNotOwner means only that no tunnel of the node is held here. Any other // failure is returned as itself, wrapped: a session that ended under a held @@ -307,17 +400,29 @@ func (t *Registry) Detach(nodeID string, lane Lane, token int64) { // A failed open does not remove the entry. Attach and Detach decide what is // held here, and an open that removed it would race with the goroutine that // owns the session and is about to detach it. -func (t *Registry) Open(ctx context.Context, nodeID string, lane Lane) (net.Conn, error) { +func (t *Registry) Open(ctx context.Context, nodeID string, lane Lane, opts ...OpenOption) (net.Conn, error) { + var o openOptions + for _, opt := range opts { + opt(&o) + } t.mu.Lock() held, ok := t.tunnels[nodeID] var sess *Session + noBulk := false if ok { sess = held.inference - if lane == LaneBulk && held.bulk != nil && !held.bulk.IsClosed() { - sess = held.bulk + if lane == LaneBulk { + if held.bulk != nil && !held.bulk.IsClosed() { + sess = held.bulk + } else if o.noFallback { + sess, noBulk = nil, true + } } } t.mu.Unlock() + if noBulk { + return nil, fmt.Errorf("opening a bulk stream to node %q: %w", nodeID, ErrNoBulkSession) + } if sess == nil { return nil, fmt.Errorf("opening a stream to node %q: %w", nodeID, ErrNotOwner) } @@ -369,20 +474,6 @@ func (t *Registry) Holds(nodeID string) bool { return ok && !held.inference.IsClosed() } -// Held returns the nodes whose tunnels this replica holds, sorted. It answers -// what this process holds, which is another question than who the table says -// owns a node. That question is Owner. -func (t *Registry) Held() []string { - t.mu.Lock() - defer t.mu.Unlock() - out := make([]string, 0, len(t.tunnels)) - for nodeID := range t.tunnels { - out = append(out, nodeID) - } - sort.Strings(out) - return out -} - // Reclaim writes a new claim for every tunnel that is still held here, and // returns how many it wrote. // diff --git a/core/services/tunnel/registry_test.go b/core/services/tunnel/registry_test.go index 719d6dab30f9..8ae68dc2c143 100644 --- a/core/services/tunnel/registry_test.go +++ b/core/services/tunnel/registry_test.go @@ -50,7 +50,7 @@ var _ = Describe("Registry", func() { Expect(err).ToNot(HaveOccurred()) Expect(id).To(Equal("replica-a")) Expect(epoch).To(Equal(token)) - Expect(registry.Held()).To(Equal([]string{"w1"})) + Expect(registry.Holds("w1")).To(BeTrue()) }) It("refuses a nil session and an unknown lane without touching the table", func() { @@ -61,7 +61,7 @@ var _ = Describe("Registry", func() { Expect(err).To(HaveOccurred()) _, err = owner("w1") Expect(err).To(MatchError(cluster.ErrNoConnection)) - Expect(registry.Held()).To(BeEmpty()) + Expect(registry.Holds("w1")).To(BeFalse()) }) It("releases the claim and drops the entry on Detach", func() { @@ -73,7 +73,7 @@ var _ = Describe("Registry", func() { _, err = owner("w1") Expect(err).To(MatchError(cluster.ErrNoConnection)) - Expect(registry.Held()).To(BeEmpty()) + Expect(registry.Holds("w1")).To(BeFalse()) }) It("ignores a Detach with a token that is not the live one", func() { @@ -91,7 +91,7 @@ var _ = Describe("Registry", func() { Expect(err).ToNot(HaveOccurred()) Expect(id).To(Equal("replica-a")) Expect(epoch).To(Equal(live)) - Expect(registry.Held()).To(Equal([]string{"w1"})) + Expect(registry.Holds("w1")).To(BeTrue()) }) It("closes the session that a new dial replaced, and keeps the new one", func() { @@ -123,7 +123,7 @@ var _ = Describe("Registry", func() { _, err := registry.Attach(canceled, "w1", tunnel.LaneInference, frontend) Expect(err).To(HaveOccurred()) - Expect(registry.Held()).To(BeEmpty()) + Expect(registry.Holds("w1")).To(BeFalse()) _, err = owner("w1") Expect(err).To(MatchError(cluster.ErrNoConnection)) }) @@ -200,7 +200,10 @@ var _ = Describe("Registry", func() { _, err = registry.Open(ctx, "w1", tunnel.LaneInference) Expect(err).To(HaveOccurred()) Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse()) - Expect(registry.Held()).To(Equal([]string{"w1"}), "a failed open must not remove the entry") + id, err := owner("w1") + Expect(err).ToNot(HaveOccurred(), "a failed open must not release the claim") + Expect(id).To(Equal("replica-a")) + Expect(registry.Disconnect("w1")).To(BeTrue(), "a failed open must not remove the entry") }) It("blames the caller and not the tunnel when the budget of the caller is spent", func() { @@ -301,6 +304,38 @@ var _ = Describe("Registry", func() { Expect(err).ToNot(HaveOccurred()) }) + It("does not fall back when the caller asks for the bulk lane only", func() { + inference, worker := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, inference) + Expect(err).ToNot(HaveOccurred()) + serveStreams(worker, echoTCP) + + _, err = registry.Open(ctx, "w1", tunnel.LaneBulk, tunnel.WithoutFallback()) + Expect(err).To(MatchError(tunnel.ErrNoBulkSession)) + Expect(err).ToNot(MatchError(tunnel.ErrNotOwner), "the node is held here; only its bulk lane is missing") + }) + + It("does not fall back when the bulk session has ended", func() { + _, bulk, _, _, _, _ := attachBoth() + Expect(bulk.Close()).To(Succeed()) + + _, err := registry.Open(ctx, "w1", tunnel.LaneBulk, tunnel.WithoutFallback()) + Expect(err).To(MatchError(tunnel.ErrNoBulkSession)) + }) + + It("still opens the bulk lane when it is there, and ignores the option on the inference lane", func() { + _, _, inferenceWorker, bulkWorker, _, _ := attachBoth() + serveStreams(inferenceWorker, echoTCP) + serveStreams(bulkWorker, echoTCP) + + st, err := registry.Open(ctx, "w1", tunnel.LaneBulk, tunnel.WithoutFallback()) + Expect(err).ToNot(HaveOccurred()) + Expect(st.Close()).To(Succeed()) + st, err = registry.Open(ctx, "w1", tunnel.LaneInference, tunnel.WithoutFallback()) + Expect(err).ToNot(HaveOccurred()) + Expect(st.Close()).To(Succeed()) + }) + It("falls back to the inference lane when the bulk session has ended", func() { _, bulk, inferenceWorker, _, _, _ := attachBoth() serveStreams(inferenceWorker, echoTCP) @@ -363,6 +398,89 @@ var _ = Describe("Registry", func() { }) }) + Describe("Disconnect", func() { + It("ends both sessions of a node, releases the claim and drops the entry", func() { + inference, bulk, inferenceWorker, bulkWorker, _, _ := func() (a, b, c, d *tunnel.Session, e, f int64) { + a, c = sessionPair(tunnel.LaneInference) + b, d = sessionPair(tunnel.LaneBulk) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, a) + Expect(err).ToNot(HaveOccurred()) + _, err = registry.Attach(ctx, "w1", tunnel.LaneBulk, b) + Expect(err).ToNot(HaveOccurred()) + return + }() + + Expect(registry.Disconnect("w1")).To(BeTrue()) + + Eventually(inference.CloseChan()).Should(BeClosed()) + Eventually(bulk.CloseChan()).Should(BeClosed()) + Eventually(inferenceWorker.CloseChan()).Should(BeClosed()) + Eventually(bulkWorker.CloseChan()).Should(BeClosed()) + _, err := owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + _, err = registry.Open(ctx, "w1", tunnel.LaneInference) + Expect(err).To(MatchError(tunnel.ErrNotOwner)) + }) + + It("says false for a node that is not held here and touches nothing", func() { + other, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w2", tunnel.LaneInference, other) + Expect(err).ToNot(HaveOccurred()) + + Expect(registry.Disconnect("w1")).To(BeFalse()) + Expect(registry.Holds("w2")).To(BeTrue()) + }) + + It("leaves the claim alone when the late Detach of the handler arrives", func() { + frontend, _ := sessionPair(tunnel.LaneInference) + token, err := registry.Attach(ctx, "w1", tunnel.LaneInference, frontend) + Expect(err).ToNot(HaveOccurred()) + Expect(registry.Disconnect("w1")).To(BeTrue()) + + // The goroutine that accepted the session detaches when it ends. + registry.Detach("w1", tunnel.LaneInference, token) + + // A worker that dialled again in between keeps its claim. + again, _ := sessionPair(tunnel.LaneInference) + live, err := registry.Attach(ctx, "w1", tunnel.LaneInference, again) + Expect(err).ToNot(HaveOccurred()) + registry.Detach("w1", tunnel.LaneInference, token) + _, epoch, err := clusterR.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(epoch).To(Equal(live)) + }) + }) + + Describe("Close", func() { + It("ends every session, releases every claim and refuses a later attach", func() { + a, aWorker := sessionPair(tunnel.LaneInference) + b, _ := sessionPair(tunnel.LaneInference) + _, err := registry.Attach(ctx, "w1", tunnel.LaneInference, a) + Expect(err).ToNot(HaveOccurred()) + _, err = registry.Attach(ctx, "w2", tunnel.LaneInference, b) + Expect(err).ToNot(HaveOccurred()) + + registry.Close() + + Eventually(aWorker.CloseChan()).Should(BeClosed()) + for _, node := range []string{"w1", "w2"} { + _, err := owner(node) + Expect(err).To(MatchError(cluster.ErrNoConnection), node) + Expect(registry.Holds(node)).To(BeFalse(), node) + } + late, _ := sessionPair(tunnel.LaneInference) + _, err = registry.Attach(ctx, "w3", tunnel.LaneInference, late) + Expect(err).To(MatchError(tunnel.ErrRegistryClosed)) + _, err = owner("w3") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("can be called twice and on a registry that holds nothing", func() { + registry.Close() + registry.Close() + }) + }) + Describe("Reclaim", func() { It("claims the held tunnels again after this replica was swept", func() { frontend, _ := sessionPair(tunnel.LaneInference) From b5adc6f9f8ed2ecce1e2f56450ff5312218764a5 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:07:17 +0000 Subject: [PATCH 048/162] fix(cluster): end the tunnel of a node that was removed or got a new credential The tunnel credential is checked when a worker dials and never again, so a session that was open stayed open after an admin deleted the node or a registration replaced its credential. Both paths now end the session that the replica holds. The replica also closes its tunnel registry before it leaves the instances table. Another replica that holds the session of the node is not told. That session ends when its worker dials again and is refused. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 10 +- core/http/app.go | 3 + core/http/endpoints/localai/nodes.go | 45 ++++++++- .../endpoints/localai/nodes_carrier_test.go | 93 ++++++++++++++++++- core/http/routes/nodes.go | 2 +- 5 files changed, 144 insertions(+), 9 deletions(-) diff --git a/core/application/distributed.go b/core/application/distributed.go index e1cc1c959992..ee3cb42ff149 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -17,7 +17,6 @@ import ( "github.com/mudler/LocalAI/core/services/agents" "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" - "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging" @@ -25,6 +24,7 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/sanitize" @@ -90,7 +90,13 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { - // First, so the peers see this replica leave before its services stop. + // The tunnels first: their claims name this replica, and a peer that + // asks who owns a worker must not be told a replica that is leaving. + if ds.Tunnels != nil { + ds.Tunnels.Close() + } + // Then the membership, so the peers see this replica leave before its + // services stop. ds.Membership.Stop() if ds.Health != nil { ds.Health.Stop() diff --git a/core/http/app.go b/core/http/app.go index 995bca4e0910..40eb11766d68 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -570,6 +570,9 @@ func API(application *application.Application) (*echo.Echo, error) { workerHTTPDial = d.WorkerHTTPDial tunnels = d.Tunnels registerOpts = append(registerOpts, localai.WithCarrierReader(d.Carriers)) + if d.Tunnels != nil { + registerOpts = append(registerOpts, localai.WithTunnelDisconnector(d.Tunnels)) + } if d.Router != nil { remoteUnloader = d.Router.Unloader() } diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index cf9085c422fd..ff5d19891943 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -131,6 +131,26 @@ type CarrierReader interface { type registerOptions struct { carriers CarrierReader + tunnels TunnelDisconnector +} + +// TunnelDisconnector ends the tunnel that this replica holds for a node. +// *tunnel.Registry is one. +type TunnelDisconnector interface { + Disconnect(nodeID string) bool +} + +// WithTunnelDisconnector makes a registration that replaces the tunnel +// credential of a node, and the removal of a node, end the tunnel session that +// the replica holds for it. The credential is checked when a worker dials, so +// without this a session that was open stays open for a node that no longer has +// the right to it. +// +// It acts on this replica only. A session held by another replica ends when its +// worker dials again and is refused, or when the session ends for another +// reason. +func WithTunnelDisconnector(d TunnelDisconnector) RegisterOption { + return func(o *registerOptions) { o.tunnels = d } } // RegisterOption changes how a registration or an approval answers. @@ -329,7 +349,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } - attachCarrierCredential(ctx, response, registry, node, natsCfg, carrierRow, haveCarrier) + attachCarrierCredential(ctx, response, registry, node, natsCfg, carrierRow, haveCarrier, opts.tunnels) if tunnelActive && expectedToken == "" && autoApprove { openRegistrationWarning.Do(func() { xlog.Warn("Worker registration is open: there is no registration token and new nodes are approved at once, " + @@ -390,13 +410,13 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr // credential for it: the tunnel token when the tunnel is active, the NATS user // JWT otherwise. A deployment with no carrier row is a deployment with NATS and // gets the answer that it always had. -func attachCarrierCredential(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, natsCfg natsauth.Config, row cluster.CarrierRow, haveCarrier bool) { +func attachCarrierCredential(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, natsCfg natsauth.Config, row cluster.CarrierRow, haveCarrier bool, tunnels TunnelDisconnector) { if haveCarrier { response["carrier"] = row.Active response["carrier_epoch"] = row.Epoch } if haveCarrier && row.Active == cluster.CarrierTunnel { - attachTunnelToken(ctx, response, registry, node) + attachTunnelToken(ctx, response, registry, node, tunnels) return } attachNatsJWT(response, node, natsCfg) @@ -428,7 +448,7 @@ func attachCarrierCredential(ctx context.Context, response map[string]any, regis // token. Registration is how a worker joins the cluster, and a tunnel problem // must not become a node that cannot join. The worker sees no tunnel_token, // says so, and tries at its next registration. -func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) { +func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, tunnels TunnelDisconnector) { if node == nil { return } @@ -436,6 +456,10 @@ func attachTunnelToken(ctx context.Context, response map[string]any, registry *n if err := registry.SetTunnelTokenHash(ctx, node.ID, ""); err != nil { xlog.Error("Failed to clear the tunnel credential of a node whose type holds none", "node", node.Name, "type", node.NodeType, "error", err) + return + } + if tunnels != nil { + tunnels.Disconnect(node.ID) } return } @@ -447,6 +471,11 @@ func attachTunnelToken(ctx context.Context, response map[string]any, registry *n return } response["tunnel_token"] = secret + // The old credential is gone from the table. A session that was opened with + // it is ended, so that the old holder has to present the new one. + if tunnels != nil && tunnels.Disconnect(node.ID) { + xlog.Info("Ended the tunnel session of a node whose credential was replaced", "node", node.Name) + } } // tunnelEligible reports whether a node of this type holds a tunnel credential. @@ -534,7 +563,8 @@ func cancelNodeLoads(ctx context.Context, registry *nodes.NodeRegistry, unloader } // DeregisterNodeEndpoint removes a backend node permanently (admin use). -func DeregisterNodeEndpoint(registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender) echo.HandlerFunc { +func DeregisterNodeEndpoint(registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) return func(c echo.Context) error { ctx := c.Request().Context() id := c.Param("id") @@ -543,6 +573,11 @@ func DeregisterNodeEndpoint(registry *nodes.NodeRegistry, unloader nodes.NodeCom xlog.Error("Failed to deregister node", "id", id, "error", err) return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to deregister node")) } + // The row is gone, and so is the right to a tunnel. Without this the + // session stays up, because the credential is checked at the dial only. + if opts.tunnels != nil { + opts.tunnels.Disconnect(id) + } return c.JSON(http.StatusOK, map[string]string{"message": "node deregistered"}) } } diff --git a/core/http/endpoints/localai/nodes_carrier_test.go b/core/http/endpoints/localai/nodes_carrier_test.go index d975d5d59351..e15d76feb33f 100644 --- a/core/http/endpoints/localai/nodes_carrier_test.go +++ b/core/http/endpoints/localai/nodes_carrier_test.go @@ -8,7 +8,9 @@ import ( "errors" "net/http" "net/http/httptest" + "slices" "strings" + "sync" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/services/cluster" @@ -248,10 +250,99 @@ var _ = Describe("Registration and the active carrier", func() { Expect(storedHash("w")).ToNot(BeEmpty()) response := map[string]any{} - attachTunnelToken(ctx, response, registry, &nodes.BackendNode{ID: id, Name: "w", NodeType: "gpu-farm"}) + attachTunnelToken(ctx, response, registry, &nodes.BackendNode{ID: id, Name: "w", NodeType: "gpu-farm"}, nil) Expect(response).ToNot(HaveKey("tunnel_token")) Expect(storedHash("w")).To(BeEmpty()) }) }) }) + +// recordingDisconnector records the nodes whose tunnel it was asked to end. +type recordingDisconnector struct { + mu sync.Mutex + nodes []string +} + +func (r *recordingDisconnector) Disconnect(nodeID string) bool { + r.mu.Lock() + defer r.mu.Unlock() + r.nodes = append(r.nodes, nodeID) + return true +} + +func (r *recordingDisconnector) seen() []string { + r.mu.Lock() + defer r.mu.Unlock() + return slices.Clone(r.nodes) +} + +var _ = Describe("Ending the tunnel session of a node that lost its right to it", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + tunnels *recordingDisconnector + natsCfg natsauth.Config + ) + tunnelRow := cluster.CarrierRow{Active: cluster.CarrierTunnel, Epoch: 7, State: cluster.StateStable} + + BeforeEach(func() { + ctx = context.Background() + var err error + registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + tunnels = &recordingDisconnector{} + }) + + register := func() string { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"name":"w"}`)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + h := RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg, + WithCarrierReader(fixedCarrier{row: tunnelRow}), WithTunnelDisconnector(tunnels)) + Expect(h(e.NewContext(req, rec))).To(Succeed()) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp["id"].(string) + } + + It("ends the session when a registration replaces the credential", func() { + id := register() + Expect(tunnels.seen()).To(Equal([]string{id})) + // A second registration replaces it again. + register() + Expect(tunnels.seen()).To(HaveLen(2)) + }) + + It("ends the session when an admin removes the node", func() { + id := register() + before := len(tunnels.seen()) + + e := echo.New() + req := httptest.NewRequest(http.MethodDelete, "/", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(DeregisterNodeEndpoint(registry, nil, WithTunnelDisconnector(tunnels))(c)).To(Succeed()) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(tunnels.seen()).To(HaveLen(before + 1)) + _, err := registry.Get(ctx, id) + Expect(err).To(HaveOccurred()) + }) + + It("does not end a session when the registration mints nothing", func() { + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"name":"w"}`)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + natsRow := cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 3, State: cluster.StateStable} + h := RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg, + WithCarrierReader(fixedCarrier{row: natsRow}), WithTunnelDisconnector(tunnels)) + Expect(h(e.NewContext(req, rec))).To(Succeed()) + Expect(tunnels.seen()).To(BeEmpty()) + }) +}) diff --git a/core/http/routes/nodes.go b/core/http/routes/nodes.go index 53a49e7f4f2f..ecf0468f0aa7 100644 --- a/core/http/routes/nodes.go +++ b/core/http/routes/nodes.go @@ -82,7 +82,7 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade admin.GET("/:id", localai.GetNodeEndpoint(registry)) admin.GET("/:id/models", localai.GetNodeModelsEndpoint(registry)) - admin.DELETE("/:id", localai.DeregisterNodeEndpoint(registry, unloader)) + admin.DELETE("/:id", localai.DeregisterNodeEndpoint(registry, unloader, options...)) admin.POST("/:id/drain", localai.DrainNodeEndpoint(registry, unloader)) admin.POST("/:id/resume", localai.ResumeNodeEndpoint(registry)) admin.POST("/:id/approve", localai.ApproveNodeEndpoint(registry, authDB, hmacSecret, natsCfg, options...)) From 27d50a3d486c3506d709f0334cae91a682500acc Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:09:18 +0000 Subject: [PATCH 049/162] fix(tunnel): bound the claim of a tunnel and release it when the caller left The claim in Attach ran under the context of the request, which has no deadline. A database that did not answer held the connect handler, its hijacked socket and the gate of the node for as long as the worker stayed connected. The claim now has its own bound of 10 seconds. A request that ended while the claim ran left a row that no entry would release, because the handler closes the session when Attach fails. The claim is released in that case. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/export_test.go | 14 +++++++++++++- core/services/tunnel/registry.go | 17 ++++++++++++++++- core/services/tunnel/registry_test.go | 27 +++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 2 deletions(-) diff --git a/core/services/tunnel/export_test.go b/core/services/tunnel/export_test.go index 611cc6b36776..e70144f8e37d 100644 --- a/core/services/tunnel/export_test.go +++ b/core/services/tunnel/export_test.go @@ -1,6 +1,18 @@ package tunnel -import "github.com/libp2p/go-yamux/v5" +import ( + "time" + + "github.com/libp2p/go-yamux/v5" +) // SessionConfigFor exposes the yamux configuration of a lane to the specs. func SessionConfigFor(lane Lane) *yamux.Config { return sessionConfig(lane) } + +// SetClaimTimeout changes the bound of a claim for one spec and returns a +// function that restores it. +func SetClaimTimeout(d time.Duration) (restore func()) { + old := claimTimeout + claimTimeout = d + return func() { claimTimeout = old } +} diff --git a/core/services/tunnel/registry.go b/core/services/tunnel/registry.go index ad87b0a43b9c..3bc6b288e47a 100644 --- a/core/services/tunnel/registry.go +++ b/core/services/tunnel/registry.go @@ -37,6 +37,12 @@ var ErrRegistryClosed = errors.New("tunnel: the registry is closed") // worker dials the bulk lane again within moments, and the caller may try again. var ErrNoBulkSession = errors.New("tunnel: the node holds no bulk session on this replica") +// claimTimeout bounds the claim in Attach. The context of the caller is the one +// of a request, and a request has no deadline. A database that does not answer +// would hold the goroutine of the connect handler, its hijacked socket and the +// gate of the node for as long as the worker stays connected. +var claimTimeout = 10 * time.Second + // releaseTimeout bounds the release in Detach. Detach runs in the goroutine that // has just seen a session end, and that goroutine must not wait for a database // that went away with the session. @@ -196,10 +202,19 @@ func (t *Registry) attachInference(ctx context.Context, nodeID string, sess *Ses epoch, err := func() (int64, error) { defer t.leaveClaim(nodeID) - epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) + claimCtx, cancelClaim := context.WithTimeout(ctx, claimTimeout) + epoch, err := t.reg.Claim(claimCtx, nodeID, t.selfID) + cancelClaim() if err != nil { return 0, err } + if err := ctx.Err(); err != nil { + // The worker left, or the request ended, while the claim ran. The + // claim is in the table and no entry will ever release it, because + // the caller closes the session when Attach fails. + t.releaseClaim(nodeID, epoch) + return 0, err + } t.mu.Lock() if t.closed { diff --git a/core/services/tunnel/registry_test.go b/core/services/tunnel/registry_test.go index 8ae68dc2c143..111f0b4b0c67 100644 --- a/core/services/tunnel/registry_test.go +++ b/core/services/tunnel/registry_test.go @@ -128,6 +128,33 @@ var _ = Describe("Registry", func() { Expect(err).To(MatchError(cluster.ErrNoConnection)) }) + It("stops waiting for a claim that the database does not answer, and holds nothing", func() { + restore := tunnel.SetClaimTimeout(300 * time.Millisecond) + DeferCleanup(restore) + // A first claim makes the row. A transaction then locks it, so the + // next claim of the node waits for the lock. + first, _ := sessionPair(tunnel.LaneInference) + token, err := registry.Attach(ctx, "w1", tunnel.LaneInference, first) + Expect(err).ToNot(HaveOccurred()) + registry.Detach("w1", tunnel.LaneInference, token) + + tx := db.Begin() + DeferCleanup(func() { _ = tx.Rollback() }) + Expect(tx.Exec("SELECT 1 FROM node_connections WHERE node_id = ? FOR UPDATE", "w1").Error).To(Succeed()) + + // The context has no deadline, as the one of a request has none. + stuck, _ := sessionPair(tunnel.LaneInference) + begun := time.Now() + _, err = registry.Attach(context.Background(), "w1", tunnel.LaneInference, stuck) + Expect(err).To(HaveOccurred()) + Expect(time.Since(begun)).To(BeNumerically("<", 5*time.Second)) + Expect(registry.Holds("w1")).To(BeFalse()) + + Expect(tx.Rollback().Error).To(Succeed()) + _, err = owner("w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), "the stuck claim must not have taken the row") + }) + It("keeps the entry whose claim the row carries when two dials race", func() { const dials = 8 var wg sync.WaitGroup From 3348aa6d2cd259542abd5f9ebf655ab62bc6d00b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:11:28 +0000 Subject: [PATCH 050/162] fix(worker): warn when the tunnel dials another host over plain ws A worker that dials a frontend over ws sends its credential and every request in the clear. The warning names the fix, an https URL. A frontend on the same machine is not a finding. The specs that listed the held tunnels use Holds now, since Held is gone. The spec that checks the credential of each dial leaves the bulk lane out, because that lane drew numbers from the same counter and made the spec fail now and then. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/cluster/connect_test.go | 18 +++++++-------- core/services/worker/tunnel.go | 24 ++++++++++++++++++++ core/services/worker/tunnel_endpoint_test.go | 23 +++++++++++++++---- 3 files changed, 52 insertions(+), 13 deletions(-) diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go index cb0f433136c6..dd54060923ea 100644 --- a/core/http/endpoints/cluster/connect_test.go +++ b/core/http/endpoints/cluster/connect_test.go @@ -152,7 +152,7 @@ var _ = Describe("Connect endpoint", func() { _, resp, err := dial("?id="+nodeID, registrationToken) Expect(err).To(HaveOccurred()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) - Expect(tunnels.Held()).To(BeEmpty()) + Expect(tunnels.Holds(nodeID)).To(BeFalse()) }) It("answers 401 to a node that has no tunnel credential, whatever it presents", func() { @@ -170,7 +170,7 @@ var _ = Describe("Connect endpoint", func() { _, resp, err := dial("?id="+pending.ID, "pending-token") Expect(err).To(HaveOccurred()) Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) - Expect(tunnels.Held()).To(BeEmpty()) + Expect(tunnels.Holds(nodeID)).To(BeFalse()) Expect(nodeReg.ApproveNode(ctx, pending.ID)).To(Succeed()) ws, _, err := dial("?id="+pending.ID, "pending-token") @@ -200,7 +200,7 @@ var _ = Describe("Connect endpoint", func() { _, resp, err := dial("?id="+nodeID, nodeToken, http.Header{"Origin": []string{"https://evil.example"}}) Expect(err).To(HaveOccurred()) Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) - Expect(tunnels.Held()).To(BeEmpty()) + Expect(tunnels.Holds(nodeID)).To(BeFalse()) }) }) @@ -212,7 +212,7 @@ var _ = Describe("Connect endpoint", func() { Expect(err).ToNot(HaveOccurred()) DeferCleanup(worker.Close) - Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) owner, _, err := clusterR.Owner(ctx, nodeID) Expect(err).ToNot(HaveOccurred()) Expect(owner).To(Equal("replica-a")) @@ -241,11 +241,11 @@ var _ = Describe("Connect endpoint", func() { Expect(err).ToNot(HaveOccurred()) worker, err := tunnel.ClientSession(ws, tunnel.LaneInference) Expect(err).ToNot(HaveOccurred()) - Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) Expect(worker.Close()).To(Succeed()) - Eventually(tunnels.Held, 5*time.Second).Should(BeEmpty()) + Eventually(func() bool { return tunnels.Holds(nodeID) }, 5*time.Second).Should(BeFalse()) // The registry drops the entry first and releases the claim after it. Eventually(func() error { _, _, err := clusterR.Owner(ctx, nodeID) @@ -258,7 +258,7 @@ var _ = Describe("Connect endpoint", func() { Expect(err).ToNot(HaveOccurred()) first, err := tunnel.ClientSession(ws1, tunnel.LaneInference) Expect(err).ToNot(HaveOccurred()) - Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) ws2, _, err := dial("?id="+nodeID, nodeToken) Expect(err).ToNot(HaveOccurred()) @@ -269,7 +269,7 @@ var _ = Describe("Connect endpoint", func() { Eventually(first.CloseChan(), 5*time.Second).Should(BeClosed()) Expect(second.IsClosed()).To(BeFalse()) // The end of the first session must not remove the second. - Consistently(tunnels.Held, time.Second).Should(Equal([]string{nodeID})) + Consistently(func() bool { return tunnels.Holds(nodeID) }, time.Second).Should(BeTrue()) _, _, err = clusterR.Owner(ctx, nodeID) Expect(err).ToNot(HaveOccurred()) }) @@ -298,7 +298,7 @@ var _ = Describe("Connect endpoint", func() { inference, err := tunnel.ClientSession(wsI, tunnel.LaneInference) Expect(err).ToNot(HaveOccurred()) DeferCleanup(inference.Close) - Eventually(tunnels.Held).Should(Equal([]string{nodeID})) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) wsB, _, err := dial("?id="+nodeID+"&lane=bulk", nodeToken) Expect(err).ToNot(HaveOccurred()) diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index ae7e77455dde..8304e95d8fdf 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -218,6 +218,11 @@ func StartTunnel(ctx context.Context, cfg TunnelConfig) (*Tunnel, error) { return nil, err } + if plaintextToRemoteHost(endpoint) { + xlog.Warn("The worker tunnel dials a frontend over plain ws, so its credential and every request that crosses it are readable on the network; "+ + "register with an https URL", "frontend", endpoint) + } + // Copied, so that the routing table cannot change under the accept loop. services := make(map[string]LocalService, len(cfg.Services)) for tag, svc := range cfg.Services { @@ -812,6 +817,25 @@ func tunnelEndpoint(frontendURL, nodeID string) (string, error) { return u.String(), nil } +// plaintextToRemoteHost reports whether a websocket URL uses ws and names a host +// that is not this machine. The credential of the node and the prompts of users +// travel in the clear on such a link. A frontend on the same host, as in a +// single-machine setup or a test, is not a finding. +func plaintextToRemoteHost(endpoint string) bool { + u, err := url.Parse(endpoint) + if err != nil || u.Scheme != "ws" { + return false + } + host := u.Hostname() + if host == "localhost" { + return false + } + if ip := net.ParseIP(host); ip != nil && ip.IsLoopback() { + return false + } + return true +} + // loopbackService routes a tagged stream to a process that listens on the // loopback interface of this worker. // diff --git a/core/services/worker/tunnel_endpoint_test.go b/core/services/worker/tunnel_endpoint_test.go index 08796481db18..cf6e0b63d617 100644 --- a/core/services/worker/tunnel_endpoint_test.go +++ b/core/services/worker/tunnel_endpoint_test.go @@ -132,7 +132,7 @@ var _ = Describe("Worker tunnel against the connect endpoint", func() { }) Eventually(t.Connected, "10s").Should(BeTrue()) Eventually(t.BulkConnected, "10s").Should(BeTrue()) - Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + Eventually(func() bool { return registry.Holds(nodeID) }, "10s").Should(BeTrue()) for _, lane := range []tunnel.Lane{tunnel.LaneInference, tunnel.LaneBulk} { st, err := registry.Open(ctx, nodeID, lane) @@ -151,7 +151,7 @@ var _ = Describe("Worker tunnel against the connect endpoint", func() { It("connects both lanes again, with a new claim, after the connections to the frontend were cut", func() { t := startWorker(frontend.URL, map[string]LocalService{}) - Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + Eventually(func() bool { return registry.Holds(nodeID) }, "10s").Should(BeTrue()) Eventually(t.BulkConnected, "10s").Should(BeTrue()) _, first, err := clusterR.Owner(ctx, nodeID) Expect(err).ToNot(HaveOccurred()) @@ -167,7 +167,7 @@ var _ = Describe("Worker tunnel against the connect endpoint", func() { Eventually(t.Connected, "30s").Should(BeTrue()) Eventually(t.BulkConnected, "30s").Should(BeTrue()) - Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + Eventually(func() bool { return registry.Holds(nodeID) }, "10s").Should(BeTrue()) _, second, err := clusterR.Owner(ctx, nodeID) Expect(err).ToNot(HaveOccurred()) Expect(second).ToNot(Equal(first), "a new session is a new claim") @@ -198,7 +198,7 @@ var _ = Describe("Worker tunnel against the connect endpoint", func() { }) Eventually(t.Connected, "10s").Should(BeTrue()) Eventually(t.BulkConnected, "10s").Should(BeTrue()) - Eventually(registry.Held, "10s").Should(Equal([]string{nodeID})) + Eventually(func() bool { return registry.Holds(nodeID) }, "10s").Should(BeTrue()) // The probe sends the request frame and the 8 bytes together, so a // call costs one round trip, as a call on an open connection does. @@ -370,3 +370,18 @@ var _ = Describe("Worker tunnel against the connect endpoint", func() { }) }) }) + +var _ = DescribeTable("Warning about a plain websocket to another host", + func(endpoint string, want bool) { + Expect(plaintextToRemoteHost(endpoint)).To(Equal(want)) + }, + Entry("ws to a name", "ws://frontend.example:8080/api/cluster/connect?id=n", true), + Entry("ws to a private address", "ws://10.1.2.3:8080/api/cluster/connect?id=n", true), + Entry("ws to an address in IPv6", "ws://[2001:db8::1]:8080/api/cluster/connect?id=n", true), + Entry("wss to a name", "wss://frontend.example/api/cluster/connect?id=n", false), + Entry("ws to localhost", "ws://localhost:8080/api/cluster/connect?id=n", false), + Entry("ws to IPv4 loopback", "ws://127.0.0.1:8080/api/cluster/connect?id=n", false), + Entry("ws to another loopback address", "ws://127.0.0.2:8080/api/cluster/connect?id=n", false), + Entry("ws to IPv6 loopback", "ws://[::1]:8080/api/cluster/connect?id=n", false), + Entry("a URL that does not parse", "ws://%zz", false), +) From ffb6e99fcb0109314327bca01f9037ad94c07e5f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:12:56 +0000 Subject: [PATCH 051/162] test(worker): run the latency benchmark of the tunnel client only when its label is asked for The specs compare the delay of a probe during a download, and a loaded runner moves the numbers by more than the margin. They ran in the default suite. They now skip unless the label filter names benchmark. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/tunnel_endpoint_test.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/core/services/worker/tunnel_endpoint_test.go b/core/services/worker/tunnel_endpoint_test.go index cf6e0b63d617..3358e27a755a 100644 --- a/core/services/worker/tunnel_endpoint_test.go +++ b/core/services/worker/tunnel_endpoint_test.go @@ -9,6 +9,7 @@ import ( "net/http" "net/http/httptest" "slices" + "strings" "sync" "time" @@ -182,6 +183,14 @@ var _ = Describe("Worker tunnel against the connect endpoint", func() { const directName = "direct: transfer and probe on two TCP connections" var results = map[string][]time.Duration{} + // The specs compare latencies, and a loaded runner moves them by more than + // the margin. They run when the label is asked for. + BeforeAll(func() { + if !strings.Contains(GinkgoLabelFilter(), "benchmark") { + Skip("a latency benchmark; run it with --label-filter=benchmark") + } + }) + // run starts a worker behind the link, and runs a download of // transferBytes on transferLane while a probe runs on the inference lane. run := func(name string, transferLane tunnel.Lane) { From 7929d0c9f0755e98b0b67357fef23ebedf686504 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:13:10 +0000 Subject: [PATCH 052/162] fix(worker): cap the quick redial of the bulk lane after a run of 409 A frontend answers 409 to a bulk dial that reached a replica which does not hold the tunnel. The worker dialled again after a short wait and did not count the attempt, so a tunnel that could never be served on the bulk lane made it dial every few hundred milliseconds for ever. After 20 answers in a row, a 409 counts as a failure and the wait grows with the same backoff as any other dial that failed. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/tunnel.go | 20 ++++++++++++++++++- core/services/worker/tunnel_test.go | 31 +++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index 8304e95d8fdf..7542d16de4b9 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -67,6 +67,14 @@ const ( // right one, so the wait is short. tunnelWrongReplicaDelay = 250 * time.Millisecond + // tunnelWrongReplicaLimit is how many answers of 409 in a row get the short + // wait. A load balancer that spreads dials over replicas puts one on the right + // one within a few tries. A run longer than this is not bad luck: the + // replica that holds the tunnel cannot be reached from here, or the bulk lane + // is not served, and a dial every few hundred milliseconds for ever helps + // nobody. After the limit a 409 counts as a failure and the wait grows. + tunnelWrongReplicaLimit = 20 + // tunnelHeaderTimeout bounds how long a stream may take to send its request // frame. Without a bound, a stream that sends nothing would hold a goroutine // and a slot of the session for as long as the tunnel lives. @@ -395,9 +403,19 @@ func (t *Tunnel) connectAndServe(ctx context.Context) error { // dial, or refuses it. Either way the worker must still serve. func (t *Tunnel) runBulk(ctx context.Context, streams *sync.WaitGroup) { attempt := 0 + conflicts := 0 for ctx.Err() == nil { start := t.now() sess, err := t.dial(ctx, tunnel.LaneBulk) + if isWrongReplica(err) { + conflicts++ + if conflicts > tunnelWrongReplicaLimit { + // Treated as any other failed dial below. + err = fmt.Errorf("the frontend answered 409 %d times in a row: %w", conflicts, err) + } + } else { + conflicts = 0 + } var delay time.Duration switch { case err == nil: @@ -414,7 +432,7 @@ func (t *Tunnel) runBulk(ctx context.Context, streams *sync.WaitGroup) { attempt++ delay = tunnelBackoffDelay(attempt) xlog.Debug("worker bulk tunnel ended", "node", t.nodeID, "retry_in", delay, "error", serveErr) - case isWrongReplica(err): + case conflicts > 0 && conflicts <= tunnelWrongReplicaLimit: // The dial landed on a replica that does not hold the inference // lane. Not a failure, and not an attempt for the wait. delay = wrongReplicaDelay() diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go index 1ce771c56e2a..b04d99e24dd7 100644 --- a/core/services/worker/tunnel_test.go +++ b/core/services/worker/tunnel_test.go @@ -786,6 +786,37 @@ var _ = Describe("Worker tunnel client", func() { Eventually(tun.BulkConnected, "10s").Should(BeTrue()) }) + It("stops dialling soon after a run of 409 and backs off like any other failure", func() { + frontend = newFakeFrontend(false) + frontend.bulkStatus.Store(http.StatusConflict) + // Unbuffered, so that the loop does not run ahead of the spec. + delays := make(chan time.Duration) + start(func(c *TunnelConfig) { + c.sleep = func(ctx context.Context, d time.Duration) error { + select { + case delays <- d: + return nil + case <-ctx.Done(): + return ctx.Err() + } + } + }) + session() + var seen []time.Duration + for range tunnelWrongReplicaLimit + 8 { + var d time.Duration + Eventually(delays, "10s").Should(Receive(&d)) + seen = append(seen, d) + } + for i, d := range seen[:tunnelWrongReplicaLimit] { + Expect(d).To(BeNumerically("<=", tunnelWrongReplicaDelay*3/2), "wait %d is still the short one", i+1) + } + // After the limit the waits follow the backoff: the eighth wait after + // it is at least half of its base times 2^7, or the maximum. + last := seen[len(seen)-1] + Expect(last).To(BeNumerically(">=", tunnelBackoffMax/2), "the waits must grow after the limit") + }) + It("connects both lanes again after the inference session ends", func() { frontend = newFakeFrontend(false) start(func(c *TunnelConfig) { c.sleep = func(ctx context.Context, _ time.Duration) error { return ctx.Err() } }) From b3f3d7610dda37f33bbc5a2a367ee3b9f468aa31 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:22:48 +0000 Subject: [PATCH 053/162] feat(workerctl): name the HTTP path of each control verb A worker that a frontend reaches over a tunnel serves its control verbs as HTTP. The path of a verb is the prefix and the name of the verb with the dots turned into slashes, from one function that the worker and the frontend both call. model.op has a path like the others. The package also holds the rules every control request follows: POST only, a body of at most 8 MiB, the line format of a streamed reply, and one answer for a path that no verb claims. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/workerctl/paths.go | 68 ++++++++++++++++++++++++ core/services/workerctl/paths_test.go | 76 +++++++++++++++++++++++++++ core/services/workerctl/request.go | 68 ++++++++++++++++++++++++ 3 files changed, 212 insertions(+) create mode 100644 core/services/workerctl/paths.go create mode 100644 core/services/workerctl/paths_test.go create mode 100644 core/services/workerctl/request.go diff --git a/core/services/workerctl/paths.go b/core/services/workerctl/paths.go new file mode 100644 index 000000000000..f0cc580c9821 --- /dev/null +++ b/core/services/workerctl/paths.go @@ -0,0 +1,68 @@ +package workerctl + +import ( + "encoding/json" + "strings" +) + +// Prefix is the path prefix under which a worker serves its control verbs over +// HTTP. Everything the frontend may command a worker to do lives below it, so +// the worker can put the whole control plane behind one authentication check. +const Prefix = "/v1/control/" + +// The control verbs. The name of a verb is the same on every carrier: NATS maps +// it onto a per-node subject, and the HTTP control plane maps it onto a path +// with PathOf. The request and reply bodies are the payloads of this package, +// so a worker that is reached over NATS and one that is reached over a tunnel +// answer with the same bytes. +const ( + VerbBackendInstall = "backend.install" + VerbBackendUpgrade = "backend.upgrade" + VerbBackendStop = "backend.stop" + VerbBackendDelete = "backend.delete" + VerbBackendList = "backend.list" + VerbModelsRunning = "models.running" + VerbModelUnload = "model.unload" + VerbModelStop = "model.stop" + VerbModelDelete = "model.delete" + VerbModelOp = "model.op" + VerbNodeStop = "node.stop" + VerbFilesEnsure = "files.ensure" + VerbFilesStage = "files.stage" + VerbFilesTemp = "files.temp" + VerbFilesListDir = "files.listdir" + VerbFilesRelease = "files.release" +) + +// AllVerbs returns every control verb once. A spec that asserts a property of +// the whole set reads it from here, so the list cannot go stale when a verb is +// added. +func AllVerbs() []string { + return []string{ + VerbBackendInstall, VerbBackendUpgrade, VerbBackendStop, VerbBackendDelete, + VerbBackendList, VerbModelsRunning, VerbModelUnload, VerbModelStop, + VerbModelDelete, VerbModelOp, VerbNodeStop, + VerbFilesEnsure, VerbFilesStage, VerbFilesTemp, VerbFilesListDir, VerbFilesRelease, + } +} + +// PathOf returns the HTTP path of a verb: "backend.install" is +// "/v1/control/backend/install". One function, used by the worker that mounts +// the verb and by the frontend that calls it, so the two cannot disagree. +func PathOf(verb string) string { + return Prefix + strings.ReplaceAll(verb, ".", "/") +} + +// Envelope is one line of a streaming control response. +// +// Exactly one of Progress and Reply is set. Zero or more Progress lines are +// followed by exactly one Reply line, and the Reply line is the last thing in +// the body. A caller that reads to the end of the body without a Reply line +// learned nothing about the work: the link broke. +type Envelope struct { + Progress json.RawMessage `json:"progress,omitempty"` + Reply json.RawMessage `json:"reply,omitempty"` +} + +// ContentTypeStream is the media type of a streaming control response. +const ContentTypeStream = "application/x-ndjson" diff --git a/core/services/workerctl/paths_test.go b/core/services/workerctl/paths_test.go new file mode 100644 index 000000000000..c708fc336d6a --- /dev/null +++ b/core/services/workerctl/paths_test.go @@ -0,0 +1,76 @@ +package workerctl_test + +import ( + "net/http" + "net/http/httptest" + "strings" + + "github.com/mudler/LocalAI/core/services/workerctl" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Control paths", func() { + It("maps every verb onto one path below the prefix, and no two verbs share one", func() { + seen := map[string]string{} + for _, verb := range workerctl.AllVerbs() { + path := workerctl.PathOf(verb) + Expect(path).To(HavePrefix(workerctl.Prefix), verb) + Expect(path).ToNot(ContainSubstring("."), verb) + Expect(seen).ToNot(HaveKey(path), "%s shares a path with %s", verb, seen[path]) + seen[path] = verb + } + }) + + It("lists each verb once", func() { + Expect(workerctl.AllVerbs()).To(HaveLen(16)) + Expect(workerctl.AllVerbs()).To(ContainElements(workerctl.VerbModelOp, workerctl.VerbFilesRelease)) + seen := map[string]bool{} + for _, v := range workerctl.AllVerbs() { + Expect(seen[v]).To(BeFalse(), v) + seen[v] = true + } + }) + + It("keeps the paths of the verbs that older frontends and workers already use", func() { + Expect(workerctl.PathOf(workerctl.VerbBackendInstall)).To(Equal("/v1/control/backend/install")) + Expect(workerctl.PathOf(workerctl.VerbModelsRunning)).To(Equal("/v1/control/models/running")) + Expect(workerctl.PathOf(workerctl.VerbModelOp)).To(Equal("/v1/control/model/op")) + Expect(workerctl.PathOf(workerctl.VerbFilesListDir)).To(Equal("/v1/control/files/listdir")) + }) +}) + +var _ = Describe("Control request rules", func() { + It("refuses anything but a POST", func() { + rec := httptest.NewRecorder() + _, ok := workerctl.ReadRequestBody(rec, httptest.NewRequest(http.MethodGet, workerctl.PathOf(workerctl.VerbNodeStop), nil)) + Expect(ok).To(BeFalse()) + Expect(rec.Code).To(Equal(http.StatusMethodNotAllowed)) + }) + + It("reads a body up to the bound and refuses a longer one with 400", func() { + rec := httptest.NewRecorder() + body, ok := workerctl.ReadRequestBody(rec, httptest.NewRequest(http.MethodPost, "/", strings.NewReader("{}"))) + Expect(ok).To(BeTrue()) + Expect(string(body)).To(Equal("{}")) + + rec = httptest.NewRecorder() + big := strings.NewReader(strings.Repeat("x", workerctl.MaxRequestBytes+1)) + _, ok = workerctl.ReadRequestBody(rec, httptest.NewRequest(http.MethodPost, "/", big)) + Expect(ok).To(BeFalse()) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + }) + + It("answers an unknown path with 404 and a body that says so, bounded", func() { + rec := httptest.NewRecorder() + path := workerctl.Prefix + strings.Repeat("é", 500) + req := httptest.NewRequest(http.MethodPost, "/", nil) + req.URL.Path = path + workerctl.WriteUnknownPath(rec, req) + Expect(rec.Code).To(Equal(http.StatusNotFound)) + Expect(rec.Body.String()).To(ContainSubstring("unknown worker control path")) + Expect(len(rec.Body.String())).To(BeNumerically("<", workerctl.MaxEchoedPathBytes+80)) + Expect(strings.ToValidUTF8(rec.Body.String(), "")).To(Equal(rec.Body.String()), "a cut inside a rune") + }) +}) diff --git a/core/services/workerctl/request.go b/core/services/workerctl/request.go new file mode 100644 index 000000000000..e0391c0a15f0 --- /dev/null +++ b/core/services/workerctl/request.go @@ -0,0 +1,68 @@ +package workerctl + +import ( + "io" + "net/http" + "unicode/utf8" +) + +// The rules every control verb enforces on a request, and the one answer every +// worker gives for a path under Prefix that it does not serve. They live beside +// the path table so that two kinds of worker, which mount verbs under the same +// prefix, cannot end up with two copies that drift. + +// MaxRequestBytes bounds a control request body. The largest real body is +// BackendInstallRequest.BackendGalleries, a serialized gallery list of a few +// hundred kilobytes. Eight megabytes is therefore not a size the protocol +// needs: it is a defence against a body that never ends. +const MaxRequestBytes = 8 << 20 + +// MaxEchoedPathBytes bounds how much of an unknown control path the 404 body +// repeats. The path is chosen by the caller and the answer exists to be read in +// a log line. +const MaxEchoedPathBytes = 128 + +// ReadRequestBody enforces what every control verb requires of a request: that +// it is a POST, and that its body is bounded. It writes the refusal itself and +// reports false when it did. +// +// A GET is refused because a control verb is a command, and a liveness probe or +// a link prefetch must not be able to stop a node. +func ReadRequestBody(w http.ResponseWriter, r *http.Request) ([]byte, bool) { + if r.Method != http.MethodPost { + http.Error(w, "control verbs are POST only", http.StatusMethodNotAllowed) + return nil, false + } + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, MaxRequestBytes)) + if err != nil { + // A body that the worker could not read is not an answer about any + // backend. 400 is what the frontend reads as "the request was + // rejected". + http.Error(w, "reading the control request body: "+err.Error(), http.StatusBadRequest) + return nil, false + } + return body, true +} + +// WriteUnknownPath answers a path under Prefix that this worker serves no verb +// on. The body says what happened, because a bare 404 through a tunnel cannot be +// told from a fault of a proxy. The frontend reads this status as "this worker +// does not serve that verb", which is a statement about the worker's version and +// not about any backend. +func WriteUnknownPath(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unknown worker control path "+truncate(r.URL.Path, MaxEchoedPathBytes), http.StatusNotFound) +} + +// truncate bounds a string that a caller chose and that is about to be echoed. +// It cuts on a rune boundary, because a cut inside a multi-byte rune travels as +// a replacement character through every log that reads it. +func truncate(s string, limit int) string { + if len(s) <= limit { + return s + } + cut := limit + for cut > 0 && !utf8.RuneStart(s[cut]) { + cut-- + } + return s[:cut] + "…" +} From 8e34dd4e8279af4fa05da48e298a9df4a9777cf8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:22:48 +0000 Subject: [PATCH 054/162] feat(nodes): mount the control plane of a worker on its file transfer server The tunnel reaches a worker through the HTTP server that the worker already runs, so the control verbs go on that server. Every path under the control prefix needs the same bearer token as the file routes. The handler can get its routes after the server started, so a worker starts the server first and registers the verbs when the supervisor exists. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../nodes/file_transfer_control_test.go | 108 ++++++++++++++++++ core/services/nodes/file_transfer_server.go | 30 ++++- 2 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 core/services/nodes/file_transfer_control_test.go diff --git a/core/services/nodes/file_transfer_control_test.go b/core/services/nodes/file_transfer_control_test.go new file mode 100644 index 000000000000..f5852200e908 --- /dev/null +++ b/core/services/nodes/file_transfer_control_test.go @@ -0,0 +1,108 @@ +package nodes + +import ( + "fmt" + "io" + "net" + "net/http" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +var _ = Describe("The control plane on the file transfer server", func() { + freeAddr := func() string { + GinkgoHelper() + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + addr := l.Addr().String() + Expect(l.Close()).To(Succeed()) + return addr + } + + start := func(token string, control http.Handler) string { + GinkgoHelper() + addr := freeAddr() + dir := GinkgoT().TempDir() + srv, err := StartFileTransferServerWithControl(addr, dir, dir, dir, token, 1<<20, nil, nil, control) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { ShutdownFileTransferServer(srv) }) + Eventually(func() error { + c, err := net.DialTimeout("tcp", addr, time.Second) + if err == nil { + _ = c.Close() + } + return err + }).Should(Succeed()) + return "http://" + addr + } + + do := func(url, token string) (int, string) { + GinkgoHelper() + req, err := http.NewRequest(http.MethodPost, url, strings.NewReader("{}")) + Expect(err).ToNot(HaveOccurred()) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + body, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(body) + } + + It("serves every path under the prefix to a caller that holds the token", func() { + control := http.NewServeMux() + control.HandleFunc(workerctl.PathOf(workerctl.VerbBackendList), func(w http.ResponseWriter, _ *http.Request) { + _, _ = fmt.Fprint(w, "listed") + }) + base := start("secret", control) + + status, body := do(base+workerctl.PathOf(workerctl.VerbBackendList), "secret") + Expect(status).To(Equal(http.StatusOK)) + Expect(body).To(Equal("listed")) + }) + + It("refuses a caller with no token or the wrong one, before any verb runs", func() { + var ran int + control := http.HandlerFunc(func(http.ResponseWriter, *http.Request) { ran++ }) + base := start("secret", control) + + for _, token := range []string{"", "wrong"} { + status, _ := do(base+workerctl.PathOf(workerctl.VerbNodeStop), token) + Expect(status).To(Equal(http.StatusUnauthorized), "token %q", token) + } + Expect(ran).To(BeZero()) + }) + + It("takes verbs that are registered after the server started", func() { + control := http.NewServeMux() + base := start("secret", control) + status, _ := do(base+workerctl.PathOf(workerctl.VerbModelOp), "secret") + Expect(status).To(Equal(http.StatusNotFound), "nothing serves the verb yet") + + control.HandleFunc(workerctl.PathOf(workerctl.VerbModelOp), func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + }) + status, _ = do(base+workerctl.PathOf(workerctl.VerbModelOp), "secret") + Expect(status).To(Equal(http.StatusNoContent)) + }) + + It("serves no control path when the worker has no control plane", func() { + base := start("secret", nil) + status, _ := do(base+workerctl.PathOf(workerctl.VerbNodeStop), "secret") + Expect(status).To(Equal(http.StatusNotFound)) + }) + + It("keeps the probes open, with a control plane mounted", func() { + base := start("secret", http.NewServeMux()) + resp, err := http.Get(base + "/healthz") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + }) +}) diff --git a/core/services/nodes/file_transfer_server.go b/core/services/nodes/file_transfer_server.go index 486753b77828..11e4c9c23b88 100644 --- a/core/services/nodes/file_transfer_server.go +++ b/core/services/nodes/file_transfer_server.go @@ -21,6 +21,7 @@ import ( "github.com/gorilla/websocket" "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/safefile" @@ -57,11 +58,21 @@ func StartFileTransferServer(addr, stagingDir, modelsDir, dataDir, token string, // StartFileTransferServerWithCapacity starts the file transfer server with a // worker-local guard for per-request ephemeral inputs. func StartFileTransferServerWithCapacity(addr, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, logStore ...*model.BackendLogStore) (*http.Server, error) { + return StartFileTransferServerWithControl(addr, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, nil, logStore...) +} + +// StartFileTransferServerWithControl is StartFileTransferServerWithCapacity plus +// the control plane of the worker. A non-nil control handler serves every path +// under workerctl.Prefix, behind the same bearer check as the file routes. The +// handler may gain its routes after the server started, as http.ServeMux does, +// so a worker can start the server before the supervisor that owns the verbs +// exists. +func StartFileTransferServerWithControl(addr, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, control http.Handler, logStore ...*model.BackendLogStore) (*http.Server, error) { listener, err := net.Listen("tcp", addr) if err != nil { return nil, fmt.Errorf("listen %s: %w", addr, err) } - return startFileTransferServer(listener, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, logStore...) + return startFileTransferServerWithControl(listener, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, control, logStore...) } // StartFileTransferServerWithListener starts the server on an existing listener. @@ -79,6 +90,10 @@ func StartFileTransferServerWithReadiness(lis net.Listener, stagingDir, modelsDi } func startFileTransferServer(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, logStore ...*model.BackendLogStore) (*http.Server, error) { + return startFileTransferServerWithControl(lis, stagingDir, modelsDir, dataDir, token, maxUploadSize, readiness, capacity, nil, logStore...) +} + +func startFileTransferServerWithControl(lis net.Listener, stagingDir, modelsDir, dataDir, token string, maxUploadSize int64, readiness *WorkerReadiness, capacity EphemeralCapacity, control http.Handler, logStore ...*model.BackendLogStore) (*http.Server, error) { if err := os.MkdirAll(stagingDir, 0750); err != nil { return nil, fmt.Errorf("creating staging dir %s: %w", stagingDir, err) } @@ -158,6 +173,19 @@ func startFileTransferServer(lis net.Listener, stagingDir, modelsDir, dataDir, t } }) + // The control plane of the worker, for a worker that a frontend reaches over + // a tunnel. The bearer check is the one of the file routes: a verb can stop a + // node, so it is never served without it. + if control != nil { + mux.HandleFunc(workerctl.Prefix, func(w http.ResponseWriter, r *http.Request) { + if !checkBearerToken(r, token) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + control.ServeHTTP(w, r) + }) + } + // Backend log endpoints (only registered when a log store is provided) var ls *model.BackendLogStore if len(logStore) > 0 && logStore[0] != nil { From 1fe59fefc4aaa562da9966762aebf2503cdc9074 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:22:48 +0000 Subject: [PATCH 055/162] feat(worker): serve the control verbs over HTTP for a worker on the tunnel A worker attached to the tunnel served no lifecycle verb, so the frontend could not install, stop or renew anything on it. It now serves every verb that the NATS server serves, with the same handlers and the same payloads: lifecycle, model operations and the file staging verbs when object storage is set. A verb answers 200 with the reply, and the worker's own refusal stays in the reply. A body that cannot be read is 400, a verb the worker does not serve is 404 with a body that says so, and a method other than POST is 405. Install and upgrade stream their progress as lines and end with one reply line. HTTP delivers requests of a verb at the same time, where NATS delivered them one by one, so a spec runs renewals, completions and the watchdog against each other under the race detector. files.ensure shares one download between callers. It now runs the download without the context of the caller that started it, so a caller that leaves does not fail the others. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/control_http.go | 214 +++++++++ core/services/worker/control_http_test.go | 503 ++++++++++++++++++++++ core/services/worker/control_server.go | 40 +- core/services/worker/file_staging.go | 7 +- core/services/worker/worker.go | 43 +- 5 files changed, 771 insertions(+), 36 deletions(-) create mode 100644 core/services/worker/control_http.go create mode 100644 core/services/worker/control_http_test.go diff --git a/core/services/worker/control_http.go b/core/services/worker/control_http.go new file mode 100644 index 000000000000..3829481371b9 --- /dev/null +++ b/core/services/worker/control_http.go @@ -0,0 +1,214 @@ +package worker + +import ( + "encoding/json" + "fmt" + "net/http" + "slices" + "sync" + + "github.com/mudler/LocalAI/core/services/workerctl" + "github.com/mudler/xlog" +) + +// httpControlServer serves the control verbs of this worker as HTTP, under +// workerctl.Prefix. The worker mounts it on the HTTP server it already runs, and +// the frontend reaches that server through the tunnel with the http stream tag. +// +// It is the second carrier of the verbs that natsControlServer serves. The +// handlers are the same functions and the bodies are the same payloads of +// package workerctl, so a worker that is reached over NATS and one that is +// reached over a tunnel answer with the same bytes. +// +// The answer of a verb is a 200 with the reply, and the reply carries the +// worker's own refusal in its Error field. A non-2xx status is the worker +// failing to read or route the request, and it is never evidence about a +// backend: 400 for a body that cannot be read, 404 for a verb this worker does +// not serve, 405 for a method other than POST. The frontend maps each of them +// to a different condition, so a handler must not use them for its own result. +type httpControlServer struct { + mux *http.ServeMux + + mu sync.Mutex + served map[controlVerb]bool +} + +func newHTTPControlServer() *httpControlServer { + s := &httpControlServer{mux: http.NewServeMux(), served: map[controlVerb]bool{}} + // Every path under the prefix that no verb claims. The body says what + // happened, because a bare 404 through a tunnel cannot be told from a fault of + // a proxy. + s.mux.HandleFunc(workerctl.Prefix, workerctl.WriteUnknownPath) + return s +} + +// ServeHTTP implements http.Handler. +func (s *httpControlServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + s.mux.ServeHTTP(w, r) +} + +// route claims the path of a verb. A verb outside the vocabulary or one claimed +// twice is an error that names the verb: the NATS server refuses a verb it has +// no subject for in the same way. +func (s *httpControlServer) route(v controlVerb) (string, error) { + if !slices.Contains(workerctl.AllVerbs(), string(v)) { + return "", fmt.Errorf("serving %s: no HTTP path for that control verb", v) + } + s.mu.Lock() + defer s.mu.Unlock() + if s.served[v] { + return "", fmt.Errorf("serving %s: the verb is already served", v) + } + s.served[v] = true + return workerctl.PathOf(string(v)), nil +} + +// handle serves a verb that answers with one JSON body, or with 204 when the +// handler returns no reply. Requests of one verb are delivered concurrently: +// the NATS server delivers them one at a time, and the handlers are written to +// be called from several goroutines. +func (s *httpControlServer) handle(v controlVerb, h controlHandler) error { + path, err := s.route(v) + if err != nil { + return err + } + s.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + body, ok := workerctl.ReadRequestBody(w, r) + if !ok { + return + } + reply, undecodable := h(r.Context(), body) + if undecodable != nil { + // The body is not a request of this verb. The typed refusal in reply + // is what the NATS server sends. Here the request failed before it + // was a request, and 400 says that without making it an answer of the + // worker about any backend. + http.Error(w, fmt.Sprintf("invalid %s request: %v", v, undecodable), http.StatusBadRequest) + return + } + if reply == nil { + w.WriteHeader(http.StatusNoContent) + return + } + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(reply); err != nil { + xlog.Debug("worker control reply could not be written", "verb", v, "error", err) + } + }) + return nil +} + +// handleWithProgress serves a verb that runs for a long time and reports +// progress while it runs. The body is a stream of workerctl.Envelope lines: any +// number of progress lines and then exactly one reply line. +func (s *httpControlServer) handleWithProgress(v controlVerb, h progressControlHandler) error { + path, err := s.route(v) + if err != nil { + return err + } + s.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + body, ok := workerctl.ReadRequestBody(w, r) + if !ok { + return + } + stream := newNDJSONStream(w) + reply, undecodable := h(r.Context(), body, stream.progress) + if undecodable != nil { + // Nothing was written yet unless a handler reported progress before it + // decoded, which none does. The status is still open. + stream.abort(fmt.Sprintf("invalid %s request: %v", v, undecodable)) + return + } + stream.reply(reply) + }) + return nil +} + +// ndjsonStream writes the lines of one streaming control response. +// +// The mutex is not optional. A progress line can come from the timer goroutine +// of a debounce, so without it a progress write could interleave with the reply +// and put a torn line on the wire. done keeps the contract that the reply is the +// last line: a progress event that arrives after it is dropped. +type ndjsonStream struct { + mu sync.Mutex + w http.ResponseWriter + enc *json.Encoder + done bool + started bool +} + +func newNDJSONStream(w http.ResponseWriter) *ndjsonStream { + return &ndjsonStream{w: w, enc: json.NewEncoder(w)} +} + +// start writes the head of the response. It is delayed to the first line so +// that a request that fails before any line can still answer with a status. +func (n *ndjsonStream) start() { + if n.started { + return + } + n.started = true + n.w.Header().Set("Content-Type", workerctl.ContentTypeStream) + // The caller reads line by line, and the first line can come minutes before + // the last. + n.w.Header().Set("X-Content-Type-Options", "nosniff") + n.w.WriteHeader(http.StatusOK) +} + +func (n *ndjsonStream) progress(ev workerctl.BackendInstallProgressEvent) { + raw, err := json.Marshal(ev) + if err != nil { + xlog.Debug("worker control progress event could not be encoded", "error", err) + return + } + n.mu.Lock() + defer n.mu.Unlock() + if n.done { + return + } + n.start() + n.write(workerctl.Envelope{Progress: raw}) +} + +// reply writes the one reply line and closes the stream to later progress. +func (n *ndjsonStream) reply(v any) { + raw, err := json.Marshal(v) + if err != nil { + // The caller waits for a reply line, and a body that ends without one + // cannot be told from a link that broke. A reply that says the worker + // could not encode its own answer can. + xlog.Error("worker control reply could not be encoded", "error", err) + raw = json.RawMessage(`{"success":false,"error":"the worker could not encode its own reply"}`) + } + n.mu.Lock() + defer n.mu.Unlock() + if n.done { + return + } + n.start() + n.write(workerctl.Envelope{Reply: raw}) + n.done = true +} + +// abort answers a request that failed before it started to stream. +func (n *ndjsonStream) abort(msg string) { + n.mu.Lock() + defer n.mu.Unlock() + if n.started || n.done { + return + } + n.done = true + http.Error(n.w, msg, http.StatusBadRequest) +} + +// write encodes one line and flushes it. The caller holds n.mu. +func (n *ndjsonStream) write(env workerctl.Envelope) { + if err := n.enc.Encode(env); err != nil { + xlog.Debug("worker control stream line could not be written", "error", err) + return + } + if f, ok := n.w.(http.Flusher); ok { + f.Flush() + } +} diff --git a/core/services/worker/control_http_test.go b/core/services/worker/control_http_test.go new file mode 100644 index 000000000000..01b5c0f7d0ab --- /dev/null +++ b/core/services/worker/control_http_test.go @@ -0,0 +1,503 @@ +package worker + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "go/ast" + "go/parser" + "go/token" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "slices" + "strconv" + "strings" + "sync" + "syscall" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// natsSubjectOf maps a verb onto the subject the NATS server uses for it, so a +// spec can ask both carriers the same question. +func natsSubjectOf(v string, node string) string { + subject, err := newNATSControlServer(newRecordingBus(), node).subject(controlVerb(v)) + Expect(err).ToNot(HaveOccurred(), v) + return subject +} + +var _ = Describe("Worker control verbs over HTTP", func() { + var ( + sigCh chan os.Signal + s *backendSupervisor + ctl *httpControlServer + srv *httptest.Server + ) + + BeforeEach(func() { + sigCh = make(chan os.Signal, 1) + s = newLifecycleTestSupervisor(sigCh) + ctl = newHTTPControlServer() + Expect(s.registerLifecycleVerbs(ctl)).To(Succeed()) + srv = httptest.NewServer(ctl) + DeferCleanup(srv.Close) + }) + + post := func(verb string, body string) (*http.Response, string) { + GinkgoHelper() + resp, err := http.Post(srv.URL+workerctl.PathOf(verb), "application/json", strings.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + raw, err := io.ReadAll(resp.Body) + Expect(err).ToNot(HaveOccurred()) + return resp, strings.TrimSpace(string(raw)) + } + + Describe("the vocabulary", func() { + It("gives every verb a path on this server and a subject on the NATS server", func() { + for _, verb := range workerctl.AllVerbs() { + fresh := newHTTPControlServer() + Expect(fresh.handle(controlVerb(verb), func(context.Context, []byte) (any, error) { return nil, nil })).To(Succeed(), verb) + Expect(natsSubjectOf(verb, "n1")).ToNot(BeEmpty(), verb) + } + }) + + It("knows every verb constant that the control server declares", func() { + fset := token.NewFileSet() + f, err := parser.ParseFile(fset, "control_server.go", nil, 0) + Expect(err).ToNot(HaveOccurred()) + var declared []string + ast.Inspect(f, func(n ast.Node) bool { + vs, ok := n.(*ast.ValueSpec) + if !ok || len(vs.Names) != 1 || !strings.HasPrefix(vs.Names[0].Name, "verb") || len(vs.Values) != 1 { + return true + } + sel, ok := vs.Values[0].(*ast.SelectorExpr) + if !ok { + return true + } + declared = append(declared, sel.Sel.Name) + return true + }) + Expect(declared).To(HaveLen(len(workerctl.AllVerbs())), + "a verb constant that is not in workerctl.AllVerbs has no path; found %v", declared) + }) + + It("refuses a verb outside the vocabulary and a verb served twice, naming the verb", func() { + fresh := newHTTPControlServer() + err := fresh.handle("backend.explode", func(context.Context, []byte) (any, error) { return nil, nil }) + Expect(err).To(MatchError(ContainSubstring("backend.explode"))) + + Expect(fresh.handle(verbBackendList, func(context.Context, []byte) (any, error) { return nil, nil })).To(Succeed()) + err = fresh.handle(verbBackendList, func(context.Context, []byte) (any, error) { return nil, nil }) + Expect(err).To(MatchError(ContainSubstring("backend.list"))) + }) + + It("serves the lifecycle verbs and nothing else when no file verb is registered", func() { + for _, verb := range []string{ + workerctl.VerbBackendInstall, workerctl.VerbBackendUpgrade, workerctl.VerbBackendStop, + workerctl.VerbBackendDelete, workerctl.VerbBackendList, workerctl.VerbModelsRunning, + workerctl.VerbModelUnload, workerctl.VerbModelStop, workerctl.VerbModelDelete, + workerctl.VerbModelOp, workerctl.VerbNodeStop, + } { + resp, _ := post(verb, malformedBody) + Expect(resp.StatusCode).ToNot(Equal(http.StatusNotFound), verb) + } + for _, verb := range []string{workerctl.VerbFilesEnsure, workerctl.VerbFilesStage, workerctl.VerbFilesTemp, workerctl.VerbFilesListDir, workerctl.VerbFilesRelease} { + resp, body := post(verb, malformedBody) + Expect(resp.StatusCode).To(Equal(http.StatusNotFound), verb) + Expect(body).To(ContainSubstring("unknown worker control path"), "the 404 must say what happened") + } + }) + }) + + Describe("the shape of an answer", func() { + It("answers a verb with 200 and its JSON reply", func() { + resp, body := post(workerctl.VerbModelsRunning, "") + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(resp.Header.Get("Content-Type")).To(Equal("application/json")) + Expect(body).To(Equal(`{"models":[],"reports_operations":true}`)) + }) + + It("answers the verb that sends no reply with 204 and signals shutdown, and never blocks on a repeat", func() { + resp, body := post(workerctl.VerbNodeStop, "") + Expect(resp.StatusCode).To(Equal(http.StatusNoContent)) + Expect(body).To(BeEmpty()) + Expect(sigCh).To(Receive(Equal(syscall.SIGTERM))) + + sigCh <- syscall.SIGINT + resp, _ = post(workerctl.VerbNodeStop, "") + Expect(resp.StatusCode).To(Equal(http.StatusNoContent)) + Expect(sigCh).To(Receive(Equal(syscall.SIGINT))) + }) + + It("gives a worker's own refusal as a 200, so it is an answer and not a fault", func() { + resp, body := post(workerctl.VerbModelOp, `{"renew":["nobody-knows-this"]}`) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + var reply workerctl.OperationReply + Expect(json.Unmarshal([]byte(body), &reply)).To(Succeed()) + Expect(reply.Unknown).To(ConsistOf("nobody-knows-this")) + }) + + It("sends the same bytes as the NATS server for the same request", func() { + bus := newRecordingBus() + natsSup := newLifecycleTestSupervisor(make(chan os.Signal, 1)) + Expect(registerLifecycleForTest(natsSup, bus)).To(Succeed()) + + for _, c := range []struct{ verb, body string }{ + {workerctl.VerbModelsRunning, ``}, + {workerctl.VerbBackendList, ``}, + {workerctl.VerbModelStop, `{"model_name":"m","process_key":"m#0","operation_id":"op"}`}, + {workerctl.VerbModelOp, `{"renew":["a"],"complete":["b"]}`}, + {workerctl.VerbModelUnload, `{"model_name":"m","address":"127.0.0.1:1"}`}, + {workerctl.VerbModelDelete, `{"model_name":"m"}`}, + {workerctl.VerbBackendStop, `{"backend":"nothing-runs"}`}, + } { + var viaNATS string + Eventually(bus.deliver(natsSubjectOf(c.verb, "n1"), []byte(c.body))).Should(Receive(&viaNATS), c.verb) + _, viaHTTP := post(c.verb, c.body) + Expect(viaHTTP).To(Equal(viaNATS), c.verb) + } + }) + }) + + Describe("a request the worker cannot read", func() { + DescribeTable("is a 400 and not an answer", + func(verb string) { + resp, body := post(verb, malformedBody) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + Expect(body).To(ContainSubstring("invalid " + verb + " request")) + }, + Entry("backend.install", workerctl.VerbBackendInstall), + Entry("backend.upgrade", workerctl.VerbBackendUpgrade), + Entry("backend.stop", workerctl.VerbBackendStop), + Entry("backend.delete", workerctl.VerbBackendDelete), + Entry("model.unload", workerctl.VerbModelUnload), + Entry("model.stop", workerctl.VerbModelStop), + Entry("model.op", workerctl.VerbModelOp), + Entry("model.delete", workerctl.VerbModelDelete), + ) + + It("still answers a verb that ignores its body", func() { + for _, verb := range []string{workerctl.VerbBackendList, workerctl.VerbModelsRunning} { + resp, _ := post(verb, malformedBody) + Expect(resp.StatusCode).To(Equal(http.StatusOK), verb) + } + }) + + It("refuses a method other than POST with 405, and does not run the verb", func() { + resp, err := http.Get(srv.URL + workerctl.PathOf(workerctl.VerbNodeStop)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusMethodNotAllowed)) + Expect(sigCh).ToNot(Receive()) + }) + + It("refuses a body above the bound with 400", func() { + big := strings.NewReader(strings.Repeat("x", workerctl.MaxRequestBytes+1)) + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbModelDelete), "application/json", big) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("answers a path under the prefix that no verb claims with 404", func() { + resp, body := post("backend.explode", "{}") + Expect(resp.StatusCode).To(Equal(http.StatusNotFound)) + Expect(body).To(ContainSubstring("unknown worker control path")) + }) + }) + + Describe("the verbs that run long", func() { + readLines := func(resp *http.Response) []workerctl.Envelope { + GinkgoHelper() + var lines []workerctl.Envelope + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 0, 64<<10), 1<<20) + for sc.Scan() { + var env workerctl.Envelope + Expect(json.Unmarshal(sc.Bytes(), &env)).To(Succeed(), sc.Text()) + lines = append(lines, env) + } + Expect(sc.Err()).ToNot(HaveOccurred()) + return lines + } + + It("streams the progress of an install and ends with exactly one reply line", func() { + s.installFn = func(_ workerctl.BackendInstallRequest, _ bool, onDownload func(string, string, string, float64)) (string, error) { + onDownload("backend.tar", "1 MB", "2 MB", 50) + onDownload("backend.tar", "2 MB", "2 MB", 100) + return "127.0.0.1:50051", nil + } + body, _ := json.Marshal(workerctl.BackendInstallRequest{Backend: "vllm", OpID: "op1"}) + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbBackendInstall), "application/json", bytes.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(resp.Header.Get("Content-Type")).To(Equal(workerctl.ContentTypeStream)) + + lines := readLines(resp) + Expect(len(lines)).To(BeNumerically(">=", 2)) + for _, l := range lines[:len(lines)-1] { + Expect(l.Reply).To(BeNil()) + var ev workerctl.BackendInstallProgressEvent + Expect(json.Unmarshal(l.Progress, &ev)).To(Succeed()) + Expect(ev.OpID).To(Equal("op1")) + } + last := lines[len(lines)-1] + Expect(last.Progress).To(BeNil()) + var reply workerctl.BackendInstallReply + Expect(json.Unmarshal(last.Reply, &reply)).To(Succeed()) + Expect(reply.Success).To(BeTrue()) + Expect(reply.Address).To(HaveSuffix(":50051")) + var progressAfterReply int + for i, l := range lines { + if l.Reply != nil && i != len(lines)-1 { + progressAfterReply++ + } + } + Expect(progressAfterReply).To(BeZero(), "the reply is the last line") + }) + + It("streams the answer of an upgrade the same way", func() { + s.upgradeFn = func(_ workerctl.BackendUpgradeRequest, _ func(string, string, string, float64)) ([]string, error) { + return []string{"m#0"}, nil + } + body, _ := json.Marshal(workerctl.BackendUpgradeRequest{Backend: "vllm"}) + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbBackendUpgrade), "application/json", bytes.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + lines := readLines(resp) + Expect(lines).ToNot(BeEmpty()) + var reply workerctl.BackendUpgradeReply + Expect(json.Unmarshal(lines[len(lines)-1].Reply, &reply)).To(Succeed()) + Expect(reply.Success).To(BeTrue()) + Expect(reply.StoppedProcessKeys).To(Equal([]string{"m#0"})) + }) + + It("does not hold up a second request while an install runs", func() { + release := make(chan struct{}) + started := make(chan struct{}) + s.installFn = func(workerctl.BackendInstallRequest, bool, func(string, string, string, float64)) (string, error) { + close(started) + <-release + return "127.0.0.1:1", nil + } + body, _ := json.Marshal(workerctl.BackendInstallRequest{Backend: "vllm"}) + done := make(chan struct{}) + go func() { + defer close(done) + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbBackendInstall), "application/json", bytes.NewReader(body)) + if err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + }() + Eventually(started).Should(BeClosed()) + + resp, _ := post(workerctl.VerbModelsRunning, "") + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + close(release) + Eventually(done).Should(BeClosed()) + }) + + It("lets an install finish when the caller goes away", func() { + finished := make(chan struct{}) + release := make(chan struct{}) + s.installFn = func(workerctl.BackendInstallRequest, bool, func(string, string, string, float64)) (string, error) { + <-release + close(finished) + return "127.0.0.1:1", nil + } + body, _ := json.Marshal(workerctl.BackendInstallRequest{Backend: "vllm"}) + ctx, cancel := context.WithCancel(context.Background()) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL+workerctl.PathOf(workerctl.VerbBackendInstall), bytes.NewReader(body)) + go func() { + resp, err := http.DefaultClient.Do(req) + if err == nil { + _ = resp.Body.Close() + } + }() + time.Sleep(100 * time.Millisecond) + cancel() + close(release) + Eventually(finished).Should(BeClosed()) + }) + }) + + Describe("handlers that several requests reach at once", func() { + It("serves renewals, completions and the watchdog without a data race", func() { + var wg sync.WaitGroup + for i := range 8 { + wg.Go(func() { + id := "op-" + strconv.Itoa(i) + for range 50 { + s.beginOperation(workerctl.BackendInstallRequest{ModelID: "m", ReplicaIndex: int32(i), OperationID: id, DeadlineMs: 60_000}) + body, _ := json.Marshal(workerctl.OperationRequest{Renew: []string{id}, Complete: []string{id}}) + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbModelOp), "application/json", bytes.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + _ = resp.Body.Close() + s.expireOperations(time.Now()) + } + }) + } + wg.Wait() + Expect(slices.Collect(func(yield func(string) bool) { + s.mu.Lock() + defer s.mu.Unlock() + for id := range s.operations { + if !yield(id) { + return + } + } + })).To(BeEmpty(), "every operation was completed") + }) + }) +}) + +var _ = Describe("Worker control verbs: both servers on one handler", func() { + // The handler that a carrier registers is the same function. A spec that ran + // only against NATS would not notice a verb that one carrier cannot serve. + It("registers the same set of verbs on both servers", func() { + bus := newRecordingBus() + s := newLifecycleTestSupervisor(make(chan os.Signal, 1)) + Expect(registerLifecycleForTest(s, bus)).To(Succeed()) + + http1 := newHTTPControlServer() + Expect(newLifecycleTestSupervisor(make(chan os.Signal, 1)).registerLifecycleVerbs(http1)).To(Succeed()) + var onHTTP []string + for v := range http1.served { + onHTTP = append(onHTTP, natsSubjectOf(string(v), "n1")) + } + Expect(bus.subscribed()).To(ConsistOf(onHTTP)) + _ = messaging.SubjectNodeStop + }) +}) + +var _ = Describe("Worker file staging verbs over HTTP", func() { + var ( + cfg *Config + cacheDir string + s3 *httptest.Server + srv *httptest.Server + getGate chan struct{} + getCount int32 + countMu sync.Mutex + ) + + BeforeEach(func() { + getGate = make(chan struct{}) + getCount = 0 + s3 = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPut: + w.WriteHeader(http.StatusOK) + case http.MethodGet: + countMu.Lock() + getCount++ + countMu.Unlock() + select { + case <-getGate: + case <-r.Context().Done(): + return + } + w.Header().Set("Content-Length", "4") + _, _ = w.Write([]byte("data")) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + DeferCleanup(s3.Close) + root := canonicalWorkerTempDir() + cfg = &Config{ + ModelsPath: filepath.Join(root, "models"), + StorageURL: s3.URL, + StorageBucket: "bucket", + StorageAccessKey: "key", + StorageSecretKey: "secret", + } + Expect(os.MkdirAll(cfg.ModelsPath, 0750)).To(Succeed()) + cacheDir = filepath.Join(root, "cache") + ctl := newHTTPControlServer() + Expect(cfg.registerFileStagingVerbs(ctl, nil)).To(Succeed()) + srv = httptest.NewServer(ctl) + DeferCleanup(srv.Close) + }) + + post := func(verb, body string) string { + GinkgoHelper() + resp, err := http.Post(srv.URL+workerctl.PathOf(verb), "application/json", strings.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusOK), verb) + raw, err := io.ReadAll(resp.Body) + Expect(err).ToNot(HaveOccurred()) + return strings.TrimSpace(string(raw)) + } + + It("serves the five file verbs", func() { + for _, verb := range []string{workerctl.VerbFilesRelease, workerctl.VerbFilesEnsure, workerctl.VerbFilesStage, workerctl.VerbFilesTemp, workerctl.VerbFilesListDir} { + resp, err := http.Post(srv.URL+workerctl.PathOf(verb), "application/json", strings.NewReader(malformedBody)) + Expect(err).ToNot(HaveOccurred()) + _ = resp.Body.Close() + Expect(resp.StatusCode).ToNot(Equal(http.StatusNotFound), verb) + } + }) + + It("stages an allowed path and refuses one outside the allowed directories", func() { + path := filepath.Join(cfg.ModelsPath, "m.gguf") + Expect(os.WriteFile(path, []byte("x"), 0640)).To(Succeed()) + Expect(post(workerctl.VerbFilesStage, fmt.Sprintf(`{"local_path":%q,"key":"models/m.gguf"}`, path))).To(Equal(`{"key":"models/m.gguf"}`)) + Expect(post(workerctl.VerbFilesStage, `{"local_path":"/etc/passwd","key":"k"}`)).To(Equal(`{"error":"path outside allowed directories"}`)) + }) + + It("allocates a temp file in the cache", func() { + got := post(workerctl.VerbFilesTemp, `{}`) + Expect(got).To(MatchRegexp(`^\{"local_path":"` + filepath.Join(cacheDir, "staging-tmp") + `/localai-staging-[0-9]+\.tmp"\}$`)) + }) + + It("lets a download that two callers share end when the first caller leaves", func() { + type result struct { + body string + err error + } + ensure := func(ctx context.Context) result { + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL+workerctl.PathOf(workerctl.VerbFilesEnsure), strings.NewReader(`{"key":"models/shared.gguf"}`)) + resp, err := http.DefaultClient.Do(req) + if err != nil { + return result{err: err} + } + defer func() { _ = resp.Body.Close() }() + raw, _ := io.ReadAll(resp.Body) + return result{body: string(raw)} + } + + firstCtx, leave := context.WithCancel(context.Background()) + first := make(chan result, 1) + go func() { first <- ensure(firstCtx) }() + Eventually(func() int32 { countMu.Lock(); defer countMu.Unlock(); return getCount }).Should(BeNumerically(">=", 1)) + + second := make(chan result, 1) + go func() { second <- ensure(context.Background()) }() + time.Sleep(200 * time.Millisecond) + + leave() + Eventually(first).Should(Receive()) + close(getGate) + + var got result + Eventually(second, "20s").Should(Receive(&got)) + Expect(got.err).ToNot(HaveOccurred()) + Expect(got.body).To(ContainSubstring(`"local_path"`), "the second caller must get the file: %s", got.body) + }) +}) diff --git a/core/services/worker/control_server.go b/core/services/worker/control_server.go index f7495eb10252..a1b068102027 100644 --- a/core/services/worker/control_server.go +++ b/core/services/worker/control_server.go @@ -8,27 +8,27 @@ import ( ) // controlVerb names one control verb independent of the carrier that delivers -// it. The NATS server maps it onto a per-node subject; a tunnel server maps it -// onto a path. +// it. The NATS server maps it onto a per-node subject; the HTTP server maps it +// onto a path. The names are those of package workerctl. type controlVerb string const ( - verbBackendInstall controlVerb = "backend.install" - verbBackendUpgrade controlVerb = "backend.upgrade" - verbBackendStop controlVerb = "backend.stop" - verbBackendDelete controlVerb = "backend.delete" - verbBackendList controlVerb = "backend.list" - verbModelsRunning controlVerb = "models.running" - verbModelUnload controlVerb = "model.unload" - verbModelStop controlVerb = "model.stop" - verbModelDelete controlVerb = "model.delete" - verbModelOp controlVerb = "model.op" - verbNodeStop controlVerb = "node.stop" - verbFilesEnsure controlVerb = "files.ensure" - verbFilesStage controlVerb = "files.stage" - verbFilesTemp controlVerb = "files.temp" - verbFilesListDir controlVerb = "files.listdir" - verbFilesRelease controlVerb = "files.release" + verbBackendInstall controlVerb = workerctl.VerbBackendInstall + verbBackendUpgrade controlVerb = workerctl.VerbBackendUpgrade + verbBackendStop controlVerb = workerctl.VerbBackendStop + verbBackendDelete controlVerb = workerctl.VerbBackendDelete + verbBackendList controlVerb = workerctl.VerbBackendList + verbModelsRunning controlVerb = workerctl.VerbModelsRunning + verbModelUnload controlVerb = workerctl.VerbModelUnload + verbModelStop controlVerb = workerctl.VerbModelStop + verbModelDelete controlVerb = workerctl.VerbModelDelete + verbModelOp controlVerb = workerctl.VerbModelOp + verbNodeStop controlVerb = workerctl.VerbNodeStop + verbFilesEnsure controlVerb = workerctl.VerbFilesEnsure + verbFilesStage controlVerb = workerctl.VerbFilesStage + verbFilesTemp controlVerb = workerctl.VerbFilesTemp + verbFilesListDir controlVerb = workerctl.VerbFilesListDir + verbFilesRelease controlVerb = workerctl.VerbFilesRelease ) // progressSink receives install progress while a long-running verb runs. @@ -50,7 +50,9 @@ type progressControlHandler func(ctx context.Context, body []byte, progress prog // registration returns once the carrier will deliver requests for the verb, or // an error that names the verb; a carrier-side refusal (a NATS permission // violation) is an error here, never a silent no-op. handle may deliver -// requests concurrently (the NATS server happens to serialise per verb). +// requests concurrently (the NATS server happens to serialise per verb; the +// HTTP server does not, so a handler must be safe to call from several +// goroutines). // handleWithProgress delivers each request on its own goroutine, because a verb // that runs for minutes must not hold up the next request of the same verb. type controlServer interface { diff --git a/core/services/worker/file_staging.go b/core/services/worker/file_staging.go index 3dc6a21dc4f1..bedef3de9b0d 100644 --- a/core/services/worker/file_staging.go +++ b/core/services/worker/file_staging.go @@ -109,7 +109,12 @@ type fileStagingVerbs struct { // ensure downloads an object storage key into the local cache. func (v *fileStagingVerbs) ensure(ctx context.Context, req workerctl.FileEnsureRequest) workerctl.FileEnsureReply { value, err, _ := v.ensureGroup.Do(req.Key, func() (any, error) { - return ensureWorkerFile(ctx, v.fm, v.capacity, req.Key) + // The download is shared by every caller of this key. It must not end + // when the caller that happened to start it goes away: the others wait + // for the same result, and a carrier with a request context (HTTP) + // cancels that context when its caller leaves. The NATS server passes a + // context that never ends. + return ensureWorkerFile(context.WithoutCancel(ctx), v.fm, v.capacity, req.Key) }) if err != nil { xlog.Error("File ensure failed", "key", req.Key, "error", err) diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index c5cd31c517e5..55414f5062b9 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "net" + "net/http" "os" "os/signal" "path/filepath" @@ -182,7 +183,18 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // means the worker has already registered with the frontend, so it is // mid-startup rather than broken. readiness := &nodes.WorkerReadiness{} - httpServer, err := nodes.StartFileTransferServerWithCapacity(httpAddr, stagingDir, cfg.ModelsPath, dataDir, cfg.RegistrationToken, config.DefaultMaxUploadSize, readiness, ephemeralCapacity, ml.BackendLogs()) + // A worker on the tunnel serves its control verbs on this server, because the + // tunnel reaches it with the http stream tag. The routes are added once the + // supervisor exists. + var ( + httpControl *httpControlServer + controlMux http.Handler + ) + if onTunnel { + httpControl = newHTTPControlServer() + controlMux = httpControl + } + httpServer, err := nodes.StartFileTransferServerWithControl(httpAddr, stagingDir, cfg.ModelsPath, dataDir, cfg.RegistrationToken, config.DefaultMaxUploadSize, readiness, ephemeralCapacity, controlMux, ml.BackendLogs()) if err != nil { return fmt.Errorf("starting HTTP file transfer server: %w", err) } @@ -310,24 +322,23 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // The watchdog stops load operations the controller no longer renews. go supervisor.runOperationWatchdog(shutdownCtx) + var control controlServer if natsClient != nil { - control := newNATSControlServer(natsClient, nodeID) - if err := supervisor.registerLifecycleVerbs(control); err != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("subscribing to worker lifecycle events: %w", err) - } + control = newNATSControlServer(natsClient, nodeID) + } else { + control = httpControl + } + if err := supervisor.registerLifecycleVerbs(control); err != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("serving the worker lifecycle verbs: %w", err) + } - // Serve the file staging verbs only when S3 is configured - if cfg.StorageURL != "" { - if err := cfg.registerFileStagingVerbs(control, ephemeralCapacity); err != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("subscribing to file staging subjects: %w", err) - } + // Serve the file staging verbs only when S3 is configured + if cfg.StorageURL != "" { + if err := cfg.registerFileStagingVerbs(control, ephemeralCapacity); err != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("serving the file staging verbs: %w", err) } - } else { - // The tunnel carries streams to the backends and to the file-transfer - // server of this worker. The lifecycle verbs have no route on it yet. - xlog.Warn("This worker is attached to the tunnel and serves no lifecycle verbs on it: backends are not installed or stopped by the frontend") } xlog.Info("Worker ready, waiting for backend.install events") From 009b5adf4017c1b20ecc6967eef6b6c12e8654e4 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:39:00 +0000 Subject: [PATCH 056/162] feat(nodes): send the control verbs over the HTTP control plane of a worker The control verbs now have a second carrier. TunnelControl sends each verb as an HTTP request that a dialer carries to the worker, and a dialer that opens a stream on the tunnel makes that a request to a worker that has no address. RemoteUnloaderAdapter and TunnelControl share one implementation of the verbs, nodeControl: the request each method builds, the timeout of each, how a reply is read and what follows a stop. A controlLink holds only how one request travels and how its failure is told apart. NATS sends the same requests with the same timeouts as before, the subjects and log lines are unchanged, and the existing specs of both pass. The HTTP link keeps the four conditions apart. A route that does not exist is ErrNoRoute only when the dialer says so. A timeout is not ErrNoRoute, and an expired budget of the caller is checked first. A missing acknowledgement of backend.stop is an error, where NATS reads it as an older worker. A 404 is a worker that does not serve the verb, and for backend.upgrade it becomes ErrNoRoute so that the fallback to the older install still works. An install that runs out of time is ErrWorkerStillInstalling, and a stream that ends before its reply line is not. The tunnel carrier joins the conformance suite of LoadOperationControl, with a scripted worker on a loopback socket. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/control_client.go | 272 ++++++++++++ core/services/nodes/control_errors_test.go | 5 +- core/services/nodes/control_nats.go | 148 ++++++- core/services/nodes/control_tunnel.go | 95 ++++ .../nodes/control_tunnel_conformance_test.go | 163 +++++++ core/services/nodes/control_tunnel_test.go | 407 ++++++++++++++++++ ...load_operation_control_conformance_test.go | 3 +- core/services/nodes/unloader.go | 334 +++++++------- 8 files changed, 1255 insertions(+), 172 deletions(-) create mode 100644 core/services/nodes/control_client.go create mode 100644 core/services/nodes/control_tunnel.go create mode 100644 core/services/nodes/control_tunnel_conformance_test.go create mode 100644 core/services/nodes/control_tunnel_test.go diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go new file mode 100644 index 000000000000..cdf83135f948 --- /dev/null +++ b/core/services/nodes/control_client.go @@ -0,0 +1,272 @@ +package nodes + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/workerctl" + "github.com/mudler/LocalAI/pkg/httpclient" +) + +// ControlClient sends the control verbs of a frontend to a worker as HTTP +// requests, over whatever dialer reaches that worker. In a tunnel deployment the +// dialer opens a stream with the http tag, so the request reaches the HTTP +// server that the worker already runs and no address of the worker is needed. +// +// There is one client for the frontend and not one for each verb, because the +// mapping from a failed request onto the four conditions of +// .agents/distributed-seams.md belongs in one place. See controlFailure. +type ControlClient struct { + // dialFor supplies the transport for one worker. It is per node because an + // http.Transport carries one DialContext, and a worker is reached through its + // own tunnel. A nil dialFor refuses every call. + dialFor WorkerNetDialerFor + token string + + // clients holds one HTTP client for each node, so that two verbs in a row + // reuse the stream that the transport already holds. ForgetNode drops one. + clientsMu sync.Mutex + clients map[string]*nodeHTTPClient +} + +// nodeHTTPClient is a client and the transport it was built on. The transport is +// kept beside the client because closing the idle streams of a node needs it, +// and what http.Client.Transport holds is whatever the shared constructor wrapped +// it in. +type nodeHTTPClient struct { + client *http.Client + transport *http.Transport +} + +// NewControlClient returns the client for the workers that dialFor can reach. It +// authenticates with the registration token of the deployment, as the file +// transfer to a worker does. +func NewControlClient(dialFor WorkerNetDialerFor, token string) *ControlClient { + return &ControlClient{dialFor: dialFor, token: token, clients: map[string]*nodeHTTPClient{}} +} + +// WorkerHTTPHost is the host of the URL of a request to the HTTP server of a +// worker. A worker that has an address gets it. A worker that holds a tunnel has +// none, and gets a name under the reserved suffix .invalid, which resolves +// nowhere. The name is never dialled: the dialer of the transport opens the +// stream, and the host exists because a request needs one and a log line should +// say which node it was for. +func WorkerHTTPHost(nodeID, addr string) string { + if addr != "" { + return addr + } + return nodeID + ".worker.invalid" +} + +// clientFor returns the HTTP client for a node, and builds it on first use. +// +// HTTP/2 is off, as it is for the file stager: its flow control stalls a long +// stream, and an install streams for minutes. No timeout is set on the client, +// because it would bound the body, and the body of an install stays open for as +// long as the install runs. What bounds a call is the context of its caller. +func (c *ControlClient) clientFor(nodeID string) (*http.Client, error) { + if c == nil || c.dialFor == nil { + return nil, fmt.Errorf("control request to node %q: no dialer is configured", nodeID) + } + c.clientsMu.Lock() + defer c.clientsMu.Unlock() + if cl, ok := c.clients[nodeID]; ok { + return cl.client, nil + } + dial := c.dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("control request to node %q: no dialer for the node", nodeID) + } + transport := &http.Transport{ + DialContext: dial, + ForceAttemptHTTP2: false, + MaxIdleConns: 10, + IdleConnTimeout: 90 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } + cl := httpclient.New(httpclient.WithTransport(transport)) + c.clients[nodeID] = &nodeHTTPClient{client: cl, transport: transport} + return cl, nil +} + +// ForgetNode drops the client of a node and closes the idle streams it holds. +// It is safe on a nil receiver. +func (c *ControlClient) ForgetNode(nodeID string) { + if c == nil { + return + } + c.clientsMu.Lock() + entry, ok := c.clients[nodeID] + delete(c.clients, nodeID) + c.clientsMu.Unlock() + if ok { + entry.transport.CloseIdleConnections() + } +} + +// Call sends one request and decodes the reply. reply may be nil for a verb that +// answers with no body. +// +// A reply with its Error field set is not an error here: the worker answered, +// and what it said is for the caller to read. Only a failure to reach the worker +// or to read its answer comes back as an error. +func (c *ControlClient) Call(ctx context.Context, nodeID, verb string, req, reply any) error { + resp, err := c.do(ctx, nodeID, verb, req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode == http.StatusNoContent || reply == nil { + // Drained, so that the transport keeps the stream for the next verb. + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxControlErrorBodyBytes)) + return nil + } + if err := json.NewDecoder(resp.Body).Decode(reply); err != nil { + return controlFailure(ctx, nodeID, fmt.Errorf("decoding the %s reply: %w", verb, err)) + } + return nil +} + +// CallStreaming sends a request whose answer is a stream of lines: zero or more +// progress lines and then one reply line. It calls onProgress for each progress +// line, in the order the worker sent them, and decodes the reply line into +// reply. onProgress may be nil. +// +// onProgress runs on the goroutine of the caller. A slow callback holds up the +// body of this request and nothing else. +func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, verb string, req, reply any, onProgress func(raw json.RawMessage)) error { + resp, err := c.do(ctx, nodeID, verb, req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + + dec := json.NewDecoder(resp.Body) + var terminal json.RawMessage + for terminal == nil { + var env workerctl.Envelope + if err := dec.Decode(&env); err != nil { + // The body ended before the reply line. That is the link breaking in + // the middle of the verb, and reading it as "the install failed" would + // be a verdict about the backend that nothing learned. + if errors.Is(err, io.EOF) { + err = fmt.Errorf("the %s stream ended before its reply line: %w", verb, io.ErrUnexpectedEOF) + } + return controlFailure(ctx, nodeID, err) + } + if env.Reply != nil { + terminal = env.Reply + break + } + if env.Progress != nil && onProgress != nil { + onProgress(env.Progress) + } + } + if reply == nil { + return nil + } + if err := json.Unmarshal(terminal, reply); err != nil { + return controlFailure(ctx, nodeID, fmt.Errorf("decoding the %s reply line: %w", verb, err)) + } + return nil +} + +// do sends the request and returns the response for a status that carries an +// answer. The caller closes the body. +func (c *ControlClient) do(ctx context.Context, nodeID, verb string, req any) (*http.Response, error) { + client, err := c.clientFor(nodeID) + if err != nil { + return nil, err + } + body, err := json.Marshal(req) + if err != nil { + return nil, controlFailure(ctx, nodeID, fmt.Errorf("encoding the %s request: %w", verb, err)) + } + path := workerctl.PathOf(verb) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://"+WorkerHTTPHost(nodeID, "")+path, bytes.NewReader(body)) + if err != nil { + return nil, controlFailure(ctx, nodeID, fmt.Errorf("building the %s request: %w", verb, err)) + } + httpReq.Header.Set("Content-Type", "application/json") + if c.token != "" { + httpReq.Header.Set("Authorization", "Bearer "+c.token) + } + + resp, err := client.Do(httpReq) + if err != nil { + return nil, controlFailure(ctx, nodeID, err) + } + switch resp.StatusCode { + case http.StatusOK, http.StatusNoContent: + return resp, nil + case http.StatusNotFound: + // The worker answered about itself and not about a backend: it is older + // than this frontend and serves no such verb. The catch-all of the control + // plane gives a body that says so, which tells it from a fault of a proxy. + _ = resp.Body.Close() + return nil, fmt.Errorf("control request %s to node %q: %w", path, nodeID, errVerbNotServed) + default: + // A verb's own failure is a 200 with the reply. A status that is not 2xx + // is the worker failing to read or route the request, and none of those + // is evidence about a backend. + detail := readErrorBody(resp) + _ = resp.Body.Close() + return nil, controlFailure(ctx, nodeID, fmt.Errorf("the %s verb answered HTTP %d: %s", verb, resp.StatusCode, detail)) + } +} + +// maxControlErrorBodyBytes bounds how much of an error body reaches a log line. +const maxControlErrorBodyBytes = 512 + +func readErrorBody(resp *http.Response) string { + raw, err := io.ReadAll(io.LimitReader(resp.Body, maxControlErrorBodyBytes)) + if err != nil { + return "" + } + return string(bytes.TrimSpace(raw)) +} + +// controlFailure names the node in the error of a control request and keeps the +// chain of the cause intact, with one rule before it: the budget of the caller is +// checked first. +// +// A cause that is ErrNoRoute stays ErrNoRoute, because the dialer put it there +// and only the dialer knows that no route exists. A cause that is a refusal of +// the worker stays what it was. Nothing else is promoted: a request that failed +// for any other reason is a plain error, which is never a route and never a +// verdict about a backend. +// +// The expired budget is checked first because a worker's refusal that arrives in +// the instant the deadline passes would otherwise read as the worker's answer, +// and nothing orders the two timers. An expiry says nothing about a worker. +func controlFailure(ctx context.Context, nodeID string, err error) error { + if err == nil { + return nil + } + if ctxErr := callerRanOut(ctx); ctxErr != nil { + return fmt.Errorf("control request to node %q: the budget of the caller ran out: %w", nodeID, ctxErr) + } + return fmt.Errorf("control request to node %q: %w", nodeID, err) +} + +// callerRanOut reports whether the budget of the caller ended an attempt. It +// answers from the wall clock as well as from the context, because the timer of +// the socket and the timer of the context are two timers and nothing orders +// them: a socket can report its timeout before ctx.Err stops returning nil. +func callerRanOut(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) { + return context.DeadlineExceeded + } + return nil +} diff --git a/core/services/nodes/control_errors_test.go b/core/services/nodes/control_errors_test.go index 7883e382ea6c..18b4ebb35fec 100644 --- a/core/services/nodes/control_errors_test.go +++ b/core/services/nodes/control_errors_test.go @@ -1,6 +1,7 @@ package nodes import ( + "context" "errors" "time" @@ -28,8 +29,8 @@ var _ = Describe("Control request error classification", func() { }) request := func() (*workerctl.BackendInstallReply, error) { - return controlRequestJSON[workerctl.BackendInstallRequest, workerctl.BackendInstallReply]( - mc, subject, workerctl.BackendInstallRequest{Backend: "b"}, time.Second) + return callVerb[workerctl.BackendInstallRequest, workerctl.BackendInstallReply]( + context.Background(), &natsLink{bus: mc}, nodeID, workerctl.VerbBackendInstall, workerctl.BackendInstallRequest{Backend: "b"}, time.Second) } It("reports a subject nobody answers as ErrNoRoute", func() { diff --git a/core/services/nodes/control_nats.go b/core/services/nodes/control_nats.go index eb6b41603dc7..d2c8dc324f8c 100644 --- a/core/services/nodes/control_nats.go +++ b/core/services/nodes/control_nats.go @@ -2,14 +2,17 @@ package nodes import ( "context" + "encoding/json" "errors" "fmt" "strings" "time" + "github.com/mudler/xlog" "github.com/nats-io/nats.go" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" ) // ErrNoRoute reports that a control command could not be delivered to a node: @@ -26,18 +29,147 @@ import ( // answered with a refusal. A node that answers is present by demonstration. var ErrNoRoute = errors.New("nodes: no route to that node") -// controlRequestJSON is messaging.RequestJSON with the carrier's failure mapped -// onto the conditions this package acts on. The carrier's own sentinel is kept -// out of the unwrap chain on purpose: it names an absence, and a consumer that -// matched on it would read absence as a fact about the node. +// natsLink carries control requests as NATS requests to the per-node subject of +// the verb. It maps the failure of the carrier onto the conditions this package +// acts on, and it keeps the carrier's own sentinel out of the unwrap chain on +// purpose: that sentinel names an absence, and a consumer that matched on it +// would read absence as a fact about the node. +type natsLink struct { + bus messaging.MessagingClient +} + +func (n *natsLink) label() string { return "NATS" } + +func (n *natsLink) describe(nodeID, verb string) string { + return "subject=" + natsSubject(nodeID, verb) +} + +// natsSubject is the per-node subject of a control verb. A verb that has no +// subject is a programming error; the empty subject then fails the request. +func natsSubject(nodeID, verb string) string { + switch verb { + case workerctl.VerbBackendInstall: + return messaging.SubjectNodeBackendInstall(nodeID) + case workerctl.VerbBackendUpgrade: + return messaging.SubjectNodeBackendUpgrade(nodeID) + case workerctl.VerbBackendStop: + return messaging.SubjectNodeBackendStop(nodeID) + case workerctl.VerbBackendDelete: + return messaging.SubjectNodeBackendDelete(nodeID) + case workerctl.VerbBackendList: + return messaging.SubjectNodeBackendList(nodeID) + case workerctl.VerbModelsRunning: + return messaging.SubjectNodeModelsRunning(nodeID) + case workerctl.VerbModelUnload: + return messaging.SubjectNodeModelUnload(nodeID) + case workerctl.VerbModelStop: + return messaging.SubjectNodeModelStop(nodeID) + case workerctl.VerbModelOp: + return messaging.SubjectNodeModelOp(nodeID) + case workerctl.VerbModelDelete: + return messaging.SubjectNodeModelDelete(nodeID) + } + return "" +} + +// request sends the request on the subject of the verb. The context is not +// used: a NATS request has its own timeout. +func (n *natsLink) request(_ context.Context, nodeID, verb string, req, reply any, timeout time.Duration) error { + subject := natsSubject(nodeID, verb) + if subject == "" { + return fmt.Errorf("no NATS subject for control verb %q", verb) + } + return controlRequest(n.bus, subject, req, reply, timeout) +} + +// controlRequest is messaging.RequestJSON with the failure of the carrier mapped +// onto the conditions this package acts on: a subject that nobody answers is +// ErrNoRoute. +func controlRequest(bus messaging.MessagingClient, subject string, req, reply any, timeout time.Duration) error { + data, err := json.Marshal(req) + if err != nil { + return fmt.Errorf("marshaling request: %w", err) + } + replyData, err := bus.Request(subject, data, timeout) + if err != nil { + err = fmt.Errorf("NATS request to %s: %w", subject, err) + if errors.Is(err, nats.ErrNoResponders) { + return fmt.Errorf("%w: %v", ErrNoRoute, err) + } + return err + } + if err := json.Unmarshal(replyData, reply); err != nil { + return fmt.Errorf("unmarshaling reply from %s: %w", subject, err) + } + return nil +} + +// controlRequestJSON is controlRequest with the types of the request and the +// reply. The carrier's own sentinel is kept out of the unwrap chain on purpose: +// it names an absence, and a consumer that matched on it would read absence as a +// fact about the node. func controlRequestJSON[Req, Reply any](bus messaging.MessagingClient, subject string, req Req, timeout time.Duration) (*Reply, error) { - reply, err := messaging.RequestJSON[Req, Reply](bus, subject, req, timeout) - if err != nil && errors.Is(err, nats.ErrNoResponders) { - return nil, fmt.Errorf("%w: %v", ErrNoRoute, err) + var reply Reply + if err := controlRequest(bus, subject, req, &reply, timeout); err != nil { + return nil, err + } + return &reply, nil +} + +// requestWithProgress subscribes to the per-operation progress subject before +// it sends the request, so that no early event is missed, and drops the +// subscription when the request is over. No subscription is made when onProgress +// is nil or opID is empty: the reconciler-driven retry path and callers that +// predate progress stay silent at no cost. +func (n *natsLink) requestWithProgress(ctx context.Context, nodeID, verb, opID string, req, reply any, timeout time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) error { + sub := n.subscribeProgress(nodeID, opID, onProgress) + err := n.request(ctx, nodeID, verb, req, reply, timeout) + if sub != nil { + if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { + xlog.Warn("Failed to unsubscribe from backend progress", "nodeID", nodeID, "verb", verb, "opID", opID, "error", unsubscribeErr) + } } - return reply, err + return err } +// subscribeProgress subscribes to the per-op backend-install progress subject +// so the master can stream per-node download ticks while a worker installs or +// upgrades. Returns nil (and subscribes to nothing) when onProgress is nil or +// opID is empty — the reconciler-driven retry path and legacy callers stay +// silent at no cost. Shared by InstallBackend, UpgradeBackend, and the legacy +// force-install fallback: an upgrade is a force-reinstall, so it reuses the +// install-progress subject rather than minting a new one (no new NATS +// permission, no new rolling-update compat surface). Caller must Unsubscribe +// the returned subscription after the request completes. +func (n *natsLink) subscribeProgress(nodeID, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) messaging.Subscription { + if onProgress == nil || opID == "" { + return nil + } + progressSubject := messaging.SubjectNodeBackendInstallProgress(nodeID, opID) + s, subErr := n.bus.Subscribe(progressSubject, func(raw []byte) { + var ev workerctl.BackendInstallProgressEvent + if err := json.Unmarshal(raw, &ev); err != nil { + xlog.Debug("malformed backend progress event", "subject", progressSubject, "error", err) + return + } + // Goroutine guard: a slow onProgress callback must not stall the NATS + // reader thread. Events spawn one goroutine each, so ordering at the + // consumer is best-effort; the worker debounces to ~250ms which dwarfs + // goroutine scheduling jitter, and its final Flush() is the terminal tick. + go onProgress(ev) + }) + if subErr != nil { + xlog.Warn("Failed to subscribe to backend progress subject; proceeding without progress streaming", + "subject", progressSubject, "error", subErr) + return nil + } + return s +} + +func (n *natsLink) waitExpired(err error) bool { return isNATSTimeout(err) } + +func (n *natsLink) acknowledgementMissing(err error) bool { return isStrictRequestTimeout(err) } + // isNATSTimeout returns true if err looks like a NATS request-reply timeout. // nats.ErrTimeout is the canonical sentinel; context.DeadlineExceeded can // also surface depending on the client's path; we accept both, plus a diff --git a/core/services/nodes/control_tunnel.go b/core/services/nodes/control_tunnel.go new file mode 100644 index 000000000000..f118593a70c9 --- /dev/null +++ b/core/services/nodes/control_tunnel.go @@ -0,0 +1,95 @@ +package nodes + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/mudler/LocalAI/core/services/workerctl" + "github.com/mudler/xlog" +) + +// TunnelControl implements NodeControl over the HTTP control plane of a worker. +// It is the second carrier of the verbs that RemoteUnloaderAdapter sends over +// NATS: the same requests, the same timeouts and the same reading of a reply, +// because both are nodeControl. Only the way a request travels differs, and with +// it the way a failure is told apart. +type TunnelControl struct { + *nodeControl +} + +// NewTunnelControl returns the sender that reaches workers through client. +// installTimeout and upgradeTimeout bound the install and upgrade verbs. +func NewTunnelControl(registry ModelLocator, client *ControlClient, installTimeout, upgradeTimeout time.Duration) *TunnelControl { + return &TunnelControl{nodeControl: &nodeControl{ + registry: registry, + link: &httpLink{client: client}, + installTimeout: installTimeout, + upgradeTimeout: upgradeTimeout, + }} +} + +// httpLink carries a control request as an HTTP request. A verb that runs for +// minutes streams its progress in the body of the same response, so nothing can +// arrive before the caller listens, and nothing is dropped. +type httpLink struct { + client *ControlClient + // trace, when set, is called with the verb and the timeout of every request. + // The specs of the carrier read what was sent through it. + trace func(verb string, timeout time.Duration) + // scale, when it is above zero, shortens every timeout by that factor. The + // specs of the carrier use it to wait out a silent worker without waiting + // minutes. trace still reports the timeout of the contract. + scale float64 +} + +func (l *httpLink) bound(timeout time.Duration) time.Duration { + if l.scale > 0 { + return time.Duration(float64(timeout) * l.scale) + } + return timeout +} + +func (l *httpLink) label() string { return "tunnel" } + +func (l *httpLink) describe(_, verb string) string { return "path=" + workerctl.PathOf(verb) } + +func (l *httpLink) request(ctx context.Context, nodeID, verb string, req, reply any, timeout time.Duration) error { + if l.trace != nil { + l.trace(verb, timeout) + } + ctx, cancel := context.WithTimeout(ctx, l.bound(timeout)) + defer cancel() + return l.client.Call(ctx, nodeID, verb, req, reply) +} + +func (l *httpLink) requestWithProgress(ctx context.Context, nodeID, verb, opID string, req, reply any, timeout time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) error { + if l.trace != nil { + l.trace(verb, timeout) + } + ctx, cancel := context.WithTimeout(ctx, l.bound(timeout)) + defer cancel() + var sink func(json.RawMessage) + if onProgress != nil && opID != "" { + sink = func(raw json.RawMessage) { + var ev workerctl.BackendInstallProgressEvent + if err := json.Unmarshal(raw, &ev); err != nil { + xlog.Debug("malformed backend progress event", "node", nodeID, "error", err) + return + } + onProgress(ev) + } + } + return l.client.CallStreaming(ctx, nodeID, verb, req, reply, sink) +} + +// waitExpired is true for the timeout that this link set on the request. Nothing +// else reads as a wait that ran out: a stream that broke is not a worker that is +// still installing. +func (l *httpLink) waitExpired(err error) bool { return errors.Is(err, context.DeadlineExceeded) } + +// acknowledgementMissing is always false. An HTTP request has no way to tell a +// worker that stopped the backend and did not reply from one that did not +// receive the request, so silence is a failure here and never a success. +func (l *httpLink) acknowledgementMissing(error) bool { return false } diff --git a/core/services/nodes/control_tunnel_conformance_test.go b/core/services/nodes/control_tunnel_conformance_test.go new file mode 100644 index 000000000000..665fcdf5eddf --- /dev/null +++ b/core/services/nodes/control_tunnel_conformance_test.go @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: MIT +package nodes + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// httpLoadOperationHarness drives the HTTP carrier for the conformance of +// LoadOperationControl. The worker is a scripted HTTP server on a loopback +// socket, and the dialer is the part that stands for the tunnel: it reaches that +// server, or it fails the way the tunnel fails. +type httpLoadOperationHarness struct { + srv *httptest.Server + control *TunnelControl + link *httpLink + + mu sync.Mutex + mode string + dialErr error + bodies []sentRequest + traced []time.Duration +} + +func newHTTPLoadOperationHarness() loadOperationHarness { + h := &httpLoadOperationHarness{mode: "answer"} + h.srv = httptest.NewServer(http.HandlerFunc(h.serve)) + DeferCleanup(h.srv.Close) + + dialFor := func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + h.mu.Lock() + err := h.dialErr + h.mu.Unlock() + if err != nil { + return nil, err + } + var d net.Dialer + return d.DialContext(ctx, "tcp", h.srv.Listener.Addr().String()) + } + } + control := NewTunnelControl(&fakeModelLocator{}, NewControlClient(dialFor, "token"), 3*time.Minute, 15*time.Minute) + link := control.link.(*httpLink) + // A worker that says nothing is waited out at a hundredth of the timeout of + // the contract. The trace still reports the timeout of the contract. + link.scale = 0.01 + link.trace = func(_ string, timeout time.Duration) { + h.mu.Lock() + defer h.mu.Unlock() + h.traced = append(h.traced, timeout) + } + h.control, h.link = control, link + return h +} + +func (h *httpLoadOperationHarness) Control() LoadOperationControl { return h.control } + +func (h *httpLoadOperationHarness) NoRoute() { + h.mu.Lock() + defer h.mu.Unlock() + h.dialErr = fmt.Errorf("%w: the tunnel of the node is not held", ErrNoRoute) +} + +func (h *httpLoadOperationHarness) TimesOut() { h.setMode("hang") } + +func (h *httpLoadOperationHarness) WorkerRefuses() { h.setMode("refuse") } + +func (h *httpLoadOperationHarness) WorkerAnswers() { h.setMode("answer") } + +func (h *httpLoadOperationHarness) setMode(m string) { + h.mu.Lock() + defer h.mu.Unlock() + h.mode = m +} + +func (h *httpLoadOperationHarness) Sent() []sentRequest { + h.mu.Lock() + defer h.mu.Unlock() + out := make([]sentRequest, len(h.bodies)) + for i, b := range h.bodies { + out[i] = b + if i < len(h.traced) { + out[i].Timeout = h.traced[i] + } + } + return out +} + +// serve is the worker. It records the request and answers as the mode says. +func (h *httpLoadOperationHarness) serve(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + h.mu.Lock() + mode := h.mode + h.mu.Unlock() + + rec := sentRequest{} + var reply any + switch r.URL.Path { + case workerctl.PathOf(workerctl.VerbBackendInstall): + rec.Verb = "install" + _ = json.Unmarshal(body, &rec.Install) + if mode == "refuse" { + reply = workerctl.BackendInstallReply{Success: false, Error: "disk full"} + } else { + reply = workerctl.BackendInstallReply{Success: true, Address: "127.0.0.1:9001", ProcessInstance: "i", ReportsOperations: true} + } + case workerctl.PathOf(workerctl.VerbModelStop): + rec.Verb = "stop" + _ = json.Unmarshal(body, &rec.Stop) + if mode == "refuse" { + reply = workerctl.ModelStopReply{Matched: true, Error: "does not belong to operation"} + } else { + reply = workerctl.ModelStopReply{Matched: true, Terminated: true} + } + case workerctl.PathOf(workerctl.VerbModelOp): + rec.Verb = "op" + _ = json.Unmarshal(body, &rec.Op) + if mode == "refuse" { + reply = workerctl.OperationReply{Unknown: []string{"op"}} + } else { + reply = workerctl.OperationReply{Renewed: []string{"op"}, Completed: []string{"op"}} + } + case workerctl.PathOf(workerctl.VerbModelUnload): + rec.Verb = "unload" + _ = json.Unmarshal(body, &rec.Unload) + if mode == "refuse" { + reply = workerctl.ModelUnloadReply{Success: false, Error: "process was replaced during unload"} + } else { + reply = workerctl.ModelUnloadReply{Success: true} + } + default: + workerctl.WriteUnknownPath(w, r) + return + } + h.mu.Lock() + h.bodies = append(h.bodies, rec) + h.mu.Unlock() + + if mode == "hang" { + <-r.Context().Done() + return + } + if rec.Verb == "install" { + raw, _ := json.Marshal(reply) + w.Header().Set("Content-Type", workerctl.ContentTypeStream) + _, _ = fmt.Fprintf(w, "{\"reply\":%s}\n", strings.TrimSpace(string(raw))) + return + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(reply) +} diff --git a/core/services/nodes/control_tunnel_test.go b/core/services/nodes/control_tunnel_test.go new file mode 100644 index 000000000000..18811b257a70 --- /dev/null +++ b/core/services/nodes/control_tunnel_test.go @@ -0,0 +1,407 @@ +package nodes + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// controlWorker is a worker whose control plane a spec scripts, reached through +// a dialer that a spec can break. +type controlWorker struct { + srv *httptest.Server + + mu sync.Mutex + handlers map[string]http.HandlerFunc + dials atomic.Int32 + headers []http.Header + hosts []string + dialErr func(ctx context.Context) error +} + +func newControlWorker() *controlWorker { + w := &controlWorker{handlers: map[string]http.HandlerFunc{}} + w.srv = httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + w.mu.Lock() + h := w.handlers[r.URL.Path] + w.headers = append(w.headers, r.Header.Clone()) + w.hosts = append(w.hosts, r.Host) + w.mu.Unlock() + // The body is read first, as the real worker does: net/http watches the + // connection for a hang-up only after the handler has read the body, and a + // handler that waits for the request context needs that. + body, _ := io.ReadAll(r.Body) + r.Body = io.NopCloser(bytes.NewReader(body)) + if h == nil { + workerctl.WriteUnknownPath(rw, r) + return + } + h(rw, r) + })) + DeferCleanup(w.srv.Close) + return w +} + +func (w *controlWorker) on(verb string, h http.HandlerFunc) { + w.mu.Lock() + defer w.mu.Unlock() + w.handlers[workerctl.PathOf(verb)] = h +} + +func (w *controlWorker) replyJSON(verb string, reply any) { + w.on(verb, func(rw http.ResponseWriter, _ *http.Request) { + rw.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(rw).Encode(reply) + }) +} + +func (w *controlWorker) dialerFor() WorkerNetDialerFor { + return func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + w.dials.Add(1) + w.mu.Lock() + fail := w.dialErr + w.mu.Unlock() + if fail != nil { + if err := fail(ctx); err != nil { + return nil, err + } + } + var d net.Dialer + return d.DialContext(ctx, "tcp", w.srv.Listener.Addr().String()) + } + } +} + +func (w *controlWorker) failDial(f func(ctx context.Context) error) { + w.mu.Lock() + defer w.mu.Unlock() + w.dialErr = f +} + +var _ = Describe("The tunnel carrier of the control verbs", func() { + const node = "99999999-8888-7777-6666-555555555555" + + var ( + worker *controlWorker + control *TunnelControl + link *httpLink + traced map[string]time.Duration + tracedM sync.Mutex + locator *fakeModelLocator + ) + + BeforeEach(func() { + worker = newControlWorker() + locator = &fakeModelLocator{} + control = NewTunnelControl(locator, NewControlClient(worker.dialerFor(), "registration-token"), 3*time.Minute, 15*time.Minute) + link = control.link.(*httpLink) + traced = map[string]time.Duration{} + link.trace = func(verb string, timeout time.Duration) { + tracedM.Lock() + defer tracedM.Unlock() + traced[verb] = timeout + } + }) + + timeoutOf := func(verb string) time.Duration { + tracedM.Lock() + defer tracedM.Unlock() + return traced[verb] + } + + Describe("the timeout of each verb", func() { + It("is the one that NATS gives it", func() { + worker.replyJSON(workerctl.VerbModelStop, workerctl.ModelStopReply{Matched: true}) + worker.replyJSON(workerctl.VerbModelOp, workerctl.OperationReply{}) + worker.replyJSON(workerctl.VerbModelUnload, workerctl.ModelUnloadReply{Success: true}) + worker.replyJSON(workerctl.VerbBackendList, workerctl.BackendListReply{}) + worker.replyJSON(workerctl.VerbModelsRunning, workerctl.ModelsRunningReply{}) + worker.replyJSON(workerctl.VerbBackendStop, workerctl.BackendStopReply{Success: true}) + worker.replyJSON(workerctl.VerbBackendDelete, workerctl.BackendDeleteReply{Success: true}) + worker.replyJSON(workerctl.VerbModelDelete, workerctl.ModelDeleteReply{Success: true}) + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") }) + worker.on(workerctl.VerbBackendUpgrade, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") }) + + _, err := control.StopLoadOperation(context.Background(), node, workerctl.ModelStopRequest{ProcessKey: "m#0", OperationID: "op"}) + Expect(err).ToNot(HaveOccurred()) + _, err = control.OperationControl(node, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(control.UnloadReplica(node, NodeModel{ModelName: "m", Address: "127.0.0.1:1"})).To(Succeed()) + _, err = control.ListBackends(node) + Expect(err).ToNot(HaveOccurred()) + _, err = control.ListRunningModels(node) + Expect(err).ToNot(HaveOccurred()) + Expect(control.StopBackend(node, "b")).To(Succeed()) + _, err = control.DeleteBackend(node, "b") + Expect(err).ToNot(HaveOccurred()) + _, err = control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", nil) + Expect(err).ToNot(HaveOccurred()) + _, err = control.UpgradeBackend(node, "b", "", "", "", "", 0, "", nil) + Expect(err).ToNot(HaveOccurred()) + + Expect(timeoutOf(workerctl.VerbModelOp)).To(Equal(5*time.Second), "a renewal must fit in its cadence") + Expect(timeoutOf(workerctl.VerbModelStop)).To(Equal(10 * time.Second)) + Expect(timeoutOf(workerctl.VerbModelUnload)).To(Equal(30 * time.Second)) + Expect(timeoutOf(workerctl.VerbBackendList)).To(Equal(30*time.Second), "the last call of this verb was the list") + Expect(timeoutOf(workerctl.VerbModelsRunning)).To(Equal(10 * time.Second)) + Expect(timeoutOf(workerctl.VerbBackendStop)).To(Equal(15 * time.Second)) + Expect(timeoutOf(workerctl.VerbBackendDelete)).To(Equal(2 * time.Minute)) + Expect(timeoutOf(workerctl.VerbBackendUpgrade)).To(Equal(15 * time.Minute)) + Expect(timeoutOf(workerctl.VerbBackendInstall)).To(Equal(3 * time.Minute)) + }) + }) + + Describe("the shape of a request", func() { + It("is a POST to the path of the verb with the registration token, to a host that names the node", func() { + var method, contentType string + worker.on(workerctl.VerbModelsRunning, func(w http.ResponseWriter, r *http.Request) { + method, contentType = r.Method, r.Header.Get("Content-Type") + _, _ = io.WriteString(w, `{"models":[]}`) + }) + _, err := control.ListRunningModels(node) + Expect(err).ToNot(HaveOccurred()) + + Expect(method).To(Equal(http.MethodPost)) + Expect(contentType).To(Equal("application/json")) + worker.mu.Lock() + defer worker.mu.Unlock() + Expect(worker.headers[0].Get("Authorization")).To(Equal("Bearer registration-token")) + Expect(worker.hosts[0]).To(Equal(node + ".worker.invalid")) + }) + + It("reuses the stream of the last request, and dials again after the node is forgotten", func() { + worker.replyJSON(workerctl.VerbModelsRunning, workerctl.ModelsRunningReply{}) + for range 3 { + _, err := control.ListRunningModels(node) + Expect(err).ToNot(HaveOccurred()) + } + Expect(worker.dials.Load()).To(Equal(int32(1))) + + control.link.(*httpLink).client.ForgetNode(node) + _, err := control.ListRunningModels(node) + Expect(err).ToNot(HaveOccurred()) + Expect(worker.dials.Load()).To(Equal(int32(2))) + }) + + It("sends a nil client's calls nowhere", func() { + var none *ControlClient + none.ForgetNode(node) + err := none.Call(context.Background(), node, workerctl.VerbModelsRunning, struct{}{}, nil) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "a missing dialer is a bug of the wiring, not a route") + }) + }) + + Describe("a worker that is too old", func() { + It("turns the 404 of an upgrade into ErrNoRoute, so that the caller falls back to the install it understands", func() { + _, err := control.UpgradeBackend(node, "b", "", "", "", "", 0, "", nil) + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue(), "%v", err) + }) + + It("does not turn the 404 of another verb into ErrNoRoute", func() { + _, err := control.OperationControl(node, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "the worker answered; it is present") + Expect(errors.Is(err, errVerbNotServed)).To(BeTrue()) + }) + + It("answers PingNode from a verb that the worker serves when the other one is missing", func() { + worker.replyJSON(workerctl.VerbModelsRunning, workerctl.ModelsRunningReply{}) + Expect(control.PingNode(node)).To(Succeed()) + }) + }) + + Describe("a link that cannot carry the request", func() { + It("reports a missing tunnel as ErrNoRoute for every verb, and PingNode with it", func() { + worker.failDial(func(context.Context) error { return fmt.Errorf("%w: no tunnel is held", ErrNoRoute) }) + _, err := control.OperationControl(node, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(control.PingNode(node), ErrNoRoute)).To(BeTrue()) + _, err = control.ListBackends(node) + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue()) + }) + + It("does not report an unreachable peer as ErrNoRoute, and keeps its error for the caller", func() { + peerGone := errors.New("the replica that holds the tunnel does not answer") + worker.failDial(func(context.Context) error { return peerGone }) + _, err := control.OperationControl(node, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(errors.Is(err, peerGone)).To(BeTrue()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "an unreachable peer is not a verdict about the worker") + Expect(control.PingNode(node)).To(Succeed(), "a node is not proven gone by a peer that did not answer") + }) + + It("keeps the refusal that a worker gave while the stream was opened", func() { + refusal := errors.New("the worker refused the stream") + worker.failDial(func(context.Context) error { return refusal }) + _, err := control.ListBackends(node) + Expect(errors.Is(err, refusal)).To(BeTrue()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("gives a renewal no more than its cadence when the stream does not open", func() { + link.scale = 0.02 + worker.failDial(func(ctx context.Context) error { <-ctx.Done(); return ctx.Err() }) + begun := time.Now() + _, err := control.OperationControl(node, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "a slow tunnel is not an absent one") + Expect(time.Since(begun)).To(BeNumerically("<", time.Second)) + }) + + It("blames the caller and not the tunnel when its own deadline ends the attempt", func() { + worker.on(workerctl.VerbModelStop, func(_ http.ResponseWriter, r *http.Request) { <-r.Context().Done() }) + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + _, err := control.StopLoadOperation(ctx, node, workerctl.ModelStopRequest{ProcessKey: "m#0", OperationID: "op"}) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + }) + + Describe("a worker that fails to serve a request", func() { + DescribeTable("is not an answer about a backend, and not a route", + func(status int) { + worker.on(workerctl.VerbBackendList, func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "something the worker could not do", status) + }) + _, err := control.ListBackends(node) + Expect(err).To(MatchError(ContainSubstring(fmt.Sprintf("HTTP %d", status)))) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + Expect(errors.Is(err, errVerbNotServed)).To(BeFalse()) + }, + Entry("400, a body it could not read", http.StatusBadRequest), + Entry("401, a token it does not accept", http.StatusUnauthorized), + Entry("405, a method it refuses", http.StatusMethodNotAllowed), + Entry("500, a handler that failed", http.StatusInternalServerError), + ) + + It("does not read a reply it cannot decode as the answer of the worker", func() { + worker.on(workerctl.VerbBackendList, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"backends":`) }) + _, err := control.ListBackends(node) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + }) + + Describe("a backend stop", func() { + It("never reads silence as a stop that was done", func() { + link.scale = 0.005 + worker.on(workerctl.VerbBackendStop, func(_ http.ResponseWriter, r *http.Request) { <-r.Context().Done() }) + err := control.StopBackend(node, "b") + Expect(err).To(HaveOccurred(), "NATS may read a missing acknowledgement as an older worker; HTTP must not") + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("reports the refusal of the worker as an error with its text", func() { + worker.replyJSON(workerctl.VerbBackendStop, workerctl.BackendStopReply{Success: false, Error: "backend is busy"}) + Expect(control.StopBackend(node, "b")).To(MatchError(ContainSubstring("backend is busy"))) + }) + }) + + Describe("an install", func() { + It("hands every progress event to the caller, in order, before the reply", func() { + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { + for i := range 3 { + ev, _ := json.Marshal(workerctl.BackendInstallProgressEvent{OpID: "op1", Percentage: float64(i * 40)}) + _, _ = fmt.Fprintf(w, "{\"progress\":%s}\n", ev) + w.(http.Flusher).Flush() + } + _, _ = io.WriteString(w, `{"reply":{"success":true,"address":"127.0.0.1:9"}}`+"\n") + }) + var seen []float64 + reply, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "op1", func(ev workerctl.BackendInstallProgressEvent) { + seen = append(seen, ev.Percentage) + }) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + Expect(seen).To(Equal([]float64{0, 40, 80})) + }) + + It("sends no progress to a caller that gave no operation id", func() { + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"progress":{"op_id":"x"}}`+"\n"+`{"reply":{"success":true}}`+"\n") + }) + called := false + _, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", func(workerctl.BackendInstallProgressEvent) { called = true }) + Expect(err).ToNot(HaveOccurred()) + Expect(called).To(BeFalse()) + }) + + It("says the worker may still be installing when the wait runs out", func() { + link.scale = 0.0005 + worker.on(workerctl.VerbBackendInstall, func(_ http.ResponseWriter, r *http.Request) { <-r.Context().Done() }) + _, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", nil) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("does not say that the worker may still be installing when the stream ends without a reply", func() { + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"progress":{"op_id":"x"}}`+"\n") + }) + _, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", nil) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, io.ErrUnexpectedEOF)).To(BeTrue()) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeFalse(), "the link broke; nothing was learned about the install") + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("hands a refusal of the worker back in the reply", func() { + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"reply":{"success":false,"error":"no space left"}}`+"\n") + }) + reply, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", nil) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Error).To(Equal("no space left")) + }) + }) + + Describe("an upgrade of a backend on an older worker, after the fallback", func() { + It("sends the legacy install with Force set", func() { + var got workerctl.BackendInstallRequest + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, r *http.Request) { + _ = json.NewDecoder(r.Body).Decode(&got) + _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") + }) + _, err := control.InstallBackendForce(node, "b", "", "", "", "", 0, "op", nil) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Force).To(BeTrue()) + }) + }) + + Describe("stopping a model on the nodes that hold it", func() { + It("stops each node once, removes its replicas and reports a node it could not reach", func() { + locator.nodes = []BackendNode{{ID: node, Name: "one"}, {ID: node, Name: "dup"}} + var stops atomic.Int32 + worker.on(workerctl.VerbBackendStop, func(w http.ResponseWriter, _ *http.Request) { + stops.Add(1) + _, _ = io.WriteString(w, `{"success":true,"stopped_process_keys":["m#0"],"reports_stopped_processes":true}`) + }) + Expect(control.UnloadRemoteModel("m")).To(Succeed()) + Expect(stops.Load()).To(Equal(int32(1))) + + worker.failDial(func(context.Context) error { return fmt.Errorf("%w: gone", ErrNoRoute) }) + control.link.(*httpLink).client.ForgetNode(node) + err := control.UnloadRemoteModel("m") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue()) + }) + }) +}) diff --git a/core/services/nodes/load_operation_control_conformance_test.go b/core/services/nodes/load_operation_control_conformance_test.go index 6ac3f0dd7ce8..615681f3a7fb 100644 --- a/core/services/nodes/load_operation_control_conformance_test.go +++ b/core/services/nodes/load_operation_control_conformance_test.go @@ -45,7 +45,8 @@ type loadOperationHarness interface { } var loadOperationCarriers = map[string]func() loadOperationHarness{ - "the NATS carrier": newNATSLoadOperationHarness, + "the NATS carrier": newNATSLoadOperationHarness, + "the tunnel carrier": newHTTPLoadOperationHarness, } type natsLoadOperationHarness struct { diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 53ef79834dc1..cab5160afca8 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -2,7 +2,6 @@ package nodes import ( "context" - "encoding/json" "errors" "fmt" "time" @@ -45,19 +44,63 @@ type NodeCommandSender interface { LoadOperationControl } -// RemoteUnloaderAdapter implements NodeCommandSender and model.RemoteModelUnloader -// by publishing NATS events for backend process lifecycle. The worker process -// subscribes and handles the actual process start/stop. +// nodeControl implements NodeControl over a controlLink. It holds what every +// carrier shares: which request each method sends, the timeout of each, how a +// reply is read, and the bookkeeping that follows a stop. The link holds only how +// one request travels and how its failure is told apart. // -// This mirrors the local ModelLoader's startProcess()/deleteProcess() but -// over NATS for remote nodes. -type RemoteUnloaderAdapter struct { +// A carrier is a link. RemoteUnloaderAdapter is nodeControl over NATS, and +// TunnelControl is nodeControl over the HTTP control plane of a worker. +type nodeControl struct { registry ModelLocator - nats messaging.MessagingClient + link controlLink installTimeout time.Duration upgradeTimeout time.Duration } +// controlLink carries one control request to one node. +// +// It is also where the four conditions are told apart (see +// .agents/distributed-seams.md), because only the carrier knows what its own +// failure looks like. request returns nil error when the worker answered, and a +// refusal of the worker is then in the reply. It wraps ErrNoRoute for a route +// that does not exist right now and for nothing else. It never wraps ErrNoRoute +// for a timeout, a failure to read the reply or a worker's refusal. +type controlLink interface { + // request sends req as verb to the node and decodes the reply into reply. A + // nil reply is for a verb that answers with no body. timeout bounds the whole + // call. + request(ctx context.Context, nodeID, verb string, req, reply any, timeout time.Duration) error + // requestWithProgress is request for a verb that reports progress while it + // runs. onProgress receives the events of operation opID, and it is nil when + // the caller wants none. Every event arrives before the call returns. + requestWithProgress(ctx context.Context, nodeID, verb, opID string, req, reply any, timeout time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) error + // label names the carrier in a log line: "NATS" or "tunnel". + label() string + // describe names where a request for verb goes, for an error message. + describe(nodeID, verb string) string + // waitExpired reports whether err says that the wait for a reply ran out. + // The wait of an install or an upgrade that runs out means that the worker + // may still be working. + waitExpired(err error) bool + // acknowledgementMissing reports whether err is exactly the carrier's own + // signal that nothing replied in time, and nothing looser. A carrier that + // reads silence as an older worker that did not reply says true. A carrier + // that cannot tell silence from a failure says false. + acknowledgementMissing(err error) bool +} + +// RemoteUnloaderAdapter implements NodeControl over NATS: it sends each control +// verb as a request to a per-node subject. The worker process subscribes and +// handles the actual process start/stop. +// +// This mirrors the local ModelLoader's startProcess()/deleteProcess() but +// over NATS for remote nodes. +type RemoteUnloaderAdapter struct { + *nodeControl + nats messaging.MessagingClient +} + // NewRemoteUnloaderAdapter creates a new adapter. installTimeout and // upgradeTimeout govern the NATS request-reply deadlines for backend.install // and backend.upgrade respectively. Use @@ -65,17 +108,20 @@ type RemoteUnloaderAdapter struct { // BackendUpgradeTimeoutOrDefault() at construction. func NewRemoteUnloaderAdapter(registry ModelLocator, nats messaging.MessagingClient, installTimeout, upgradeTimeout time.Duration) *RemoteUnloaderAdapter { return &RemoteUnloaderAdapter{ - registry: registry, - nats: nats, - installTimeout: installTimeout, - upgradeTimeout: upgradeTimeout, + nodeControl: &nodeControl{ + registry: registry, + link: &natsLink{bus: nats}, + installTimeout: installTimeout, + upgradeTimeout: upgradeTimeout, + }, + nats: nats, } } // InstallTimeout returns the configured backend.install round-trip timeout. // Used by DistributedBackendManager to push NextRetryAt out by this duration // when a worker times out replying but is still installing in the background. -func (a *RemoteUnloaderAdapter) InstallTimeout() time.Duration { +func (a *nodeControl) InstallTimeout() time.Duration { return a.installTimeout } @@ -90,6 +136,7 @@ var ( _ model.RemoteModelPresenceChecker = (*RemoteUnloaderAdapter)(nil) _ ExactModelStopper = (*RemoteUnloaderAdapter)(nil) _ NodeControl = (*RemoteUnloaderAdapter)(nil) + _ NodeControl = (*TunnelControl)(nil) ) const exactModelStopTimeout = 10 * time.Second @@ -97,7 +144,7 @@ const exactModelStopTimeout = 10 * time.Second // StopModelReplica stops only the process represented by replica. Configuration // cleanup intentionally has no backend.stop fallback: an old worker that does // not understand this request leaves the quarantine row for a later retry. -func (a *RemoteUnloaderAdapter) StopModelReplica(ctx context.Context, nodeID string, replica NodeModel, force bool) (workerctl.ModelStopReply, error) { +func (a *nodeControl) StopModelReplica(ctx context.Context, nodeID string, replica NodeModel, force bool) (workerctl.ModelStopReply, error) { return a.stopModelExact(ctx, nodeID, workerctl.ModelStopRequest{ ModelName: replica.ModelName, ProcessKey: model.BackendProcessKey(replica.ModelName, replica.ReplicaIndex), @@ -113,14 +160,14 @@ func (a *RemoteUnloaderAdapter) StopModelReplica(ctx context.Context, nodeID str // worker refuses unless the operation, the process key and, when given, the // address and process instance all match its own records, so a stop of a load // that finished cannot become an unload of the model. -func (a *RemoteUnloaderAdapter) StopLoadOperation(ctx context.Context, nodeID string, req workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { +func (a *nodeControl) StopLoadOperation(ctx context.Context, nodeID string, req workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { if req.OperationID == "" { return workerctl.ModelStopReply{}, errors.New("a load stop needs an operation id") } return a.stopModelExact(ctx, nodeID, req) } -func (a *RemoteUnloaderAdapter) stopModelExact(ctx context.Context, nodeID string, req workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { +func (a *nodeControl) stopModelExact(ctx context.Context, nodeID string, req workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { if ctx == nil { ctx = context.Background() } @@ -133,7 +180,7 @@ func (a *RemoteUnloaderAdapter) stopModelExact(ctx context.Context, nodeID strin } done := make(chan result, 1) go func() { - reply, err := controlRequestJSON[workerctl.ModelStopRequest, workerctl.ModelStopReply](a.nats, messaging.SubjectNodeModelStop(nodeID), req, exactModelStopTimeout) + reply, err := callVerb[workerctl.ModelStopRequest, workerctl.ModelStopReply](ctx, a.link, nodeID, workerctl.VerbModelStop, req, exactModelStopTimeout) done <- result{reply: reply, err: err} }() @@ -151,15 +198,15 @@ func (a *RemoteUnloaderAdapter) stopModelExact(ctx context.Context, nodeID strin // OperationControl renews and completes load operations on a node. A worker // that predates operations never answers; the caller treats that as a node it // cannot confirm stops on. -func (a *RemoteUnloaderAdapter) OperationControl(nodeID string, req workerctl.OperationRequest) (*workerctl.OperationReply, error) { - return controlRequestJSON[workerctl.OperationRequest, workerctl.OperationReply](a.nats, messaging.SubjectNodeModelOp(nodeID), req, operationControlTimeout) +func (a *nodeControl) OperationControl(nodeID string, req workerctl.OperationRequest) (*workerctl.OperationReply, error) { + return callVerb[workerctl.OperationRequest, workerctl.OperationReply](context.Background(), a.link, nodeID, workerctl.VerbModelOp, req, operationControlTimeout) } const operationControlTimeout = 5 * time.Second // InstallBackendOp is InstallBackend for a load operation: the request carries // the operation id and the longest the load may run, so the worker can bound it. -func (a *RemoteUnloaderAdapter) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (a *nodeControl) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { return a.installBackend(nodeID, workerctl.BackendInstallRequest{ Backend: backendType, ModelID: modelID, @@ -176,7 +223,7 @@ func (a *RemoteUnloaderAdapter) InstallBackendOp(nodeID, backendType, modelID, g // The worker process handles a bounded Free() followed by process termination; // forced shutdown skips Free(). // This is called by ModelLoader.deleteProcess() when process == nil (remote model). -func (a *RemoteUnloaderAdapter) UnloadRemoteModel(modelName string) error { +func (a *nodeControl) UnloadRemoteModel(modelName string) error { return a.UnloadRemoteModelContext(context.Background(), modelName, false) } @@ -185,7 +232,7 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModel(modelName string) error { // nothing to stop"; ShutdownModel consults this first so it can answer 404 for // a model loaded neither locally nor anywhere in the cluster, instead of the // misleading 500 "model not found" that a local-store miss used to produce. -func (a *RemoteUnloaderAdapter) HasRemoteModel(ctx context.Context, modelName string) (bool, error) { +func (a *nodeControl) HasRemoteModel(ctx context.Context, modelName string) (bool, error) { if ctx == nil { ctx = context.Background() } @@ -198,7 +245,7 @@ func (a *RemoteUnloaderAdapter) HasRemoteModel(ctx context.Context, modelName st // UnloadRemoteModelContext is the cancellation-aware extension used by the // model loader to preserve forced shutdown across the distributed boundary. -func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error { +func (a *nodeControl) UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error { if ctx == nil { ctx = context.Background() } @@ -222,7 +269,7 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo continue } seenNodeIDs[node.ID] = struct{}{} - xlog.Info("Sending NATS backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID, "force", force) + xlog.Info("Sending "+a.link.label()+" backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID, "force", force) if err := a.stopBackend(node.ID, modelName, force); err != nil { xlog.Warn("Failed to send backend.stop", "node", node.Name, "error", err) unloadErr = errors.Join(unloadErr, fmt.Errorf("stopping model on node %s: %w", node.ID, err)) @@ -253,7 +300,7 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo // For force-reinstall (admin-driven Upgrade), use UpgradeBackend instead - // it lives on a different NATS subject so it cannot head-of-line-block // routine load traffic on the same worker. -func (a *RemoteUnloaderAdapter) InstallBackend( +func (a *nodeControl) InstallBackend( nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, @@ -271,64 +318,20 @@ func (a *RemoteUnloaderAdapter) InstallBackend( }, onProgress) } -func (a *RemoteUnloaderAdapter) installBackend(nodeID string, req workerctl.BackendInstallRequest, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { +func (a *nodeControl) installBackend(nodeID string, req workerctl.BackendInstallRequest, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { backendType, modelID, replicaIndex, opID := req.Backend, req.ModelID, int(req.ReplicaIndex), req.OpID - subject := messaging.SubjectNodeBackendInstall(nodeID) - xlog.Info("Sending NATS backend.install", "nodeID", nodeID, "backend", backendType, "modelID", modelID, "replica", replicaIndex, "opID", opID) + target := a.link.describe(nodeID, workerctl.VerbBackendInstall) + xlog.Info("Sending "+a.link.label()+" backend.install", "nodeID", nodeID, "backend", backendType, "modelID", modelID, "replica", replicaIndex, "opID", opID) - // Subscribe to the per-op progress subject BEFORE publishing the install - // request so we don't miss early events. - sub := a.subscribeProgress(nodeID, opID, onProgress) - - reply, err := controlRequestJSON[workerctl.BackendInstallRequest, workerctl.BackendInstallReply](a.nats, subject, req, a.installTimeout) - - if sub != nil { - if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { - xlog.Warn("Failed to unsubscribe from backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) - } - } - - if err != nil && isNATSTimeout(err) { - return nil, fmt.Errorf("%w (subject=%s nodeID=%s backend=%s): %v", - galleryop.ErrWorkerStillInstalling, subject, nodeID, backendType, err) + reply, err := callVerbWithProgress[workerctl.BackendInstallRequest, workerctl.BackendInstallReply]( + context.Background(), a.link, nodeID, workerctl.VerbBackendInstall, opID, req, a.installTimeout, onProgress) + if err != nil && a.link.waitExpired(err) { + return nil, fmt.Errorf("%w (%s nodeID=%s backend=%s): %v", + galleryop.ErrWorkerStillInstalling, target, nodeID, backendType, err) } return reply, err } -// subscribeProgress subscribes to the per-op backend-install progress subject -// so the master can stream per-node download ticks while a worker installs or -// upgrades. Returns nil (and subscribes to nothing) when onProgress is nil or -// opID is empty — the reconciler-driven retry path and legacy callers stay -// silent at no cost. Shared by InstallBackend, UpgradeBackend, and the legacy -// force-install fallback: an upgrade is a force-reinstall, so it reuses the -// install-progress subject rather than minting a new one (no new NATS -// permission, no new rolling-update compat surface). Caller must Unsubscribe -// the returned subscription after the request completes. -func (a *RemoteUnloaderAdapter) subscribeProgress(nodeID, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) messaging.Subscription { - if onProgress == nil || opID == "" { - return nil - } - progressSubject := messaging.SubjectNodeBackendInstallProgress(nodeID, opID) - s, subErr := a.nats.Subscribe(progressSubject, func(raw []byte) { - var ev workerctl.BackendInstallProgressEvent - if err := json.Unmarshal(raw, &ev); err != nil { - xlog.Debug("malformed backend progress event", "subject", progressSubject, "error", err) - return - } - // Goroutine guard: a slow onProgress callback must not stall the NATS - // reader thread. Events spawn one goroutine each, so ordering at the - // consumer is best-effort; the worker debounces to ~250ms which dwarfs - // goroutine scheduling jitter, and its final Flush() is the terminal tick. - go onProgress(ev) - }) - if subErr != nil { - xlog.Warn("Failed to subscribe to backend progress subject; proceeding without progress streaming", - "subject", progressSubject, "error", subErr) - return nil - } - return s -} - // UpgradeBackend sends a backend.upgrade request-reply to a worker node. // The worker stops every live process for this backend, force-reinstalls // from the gallery (overwriting the on-disk artifact), and replies. The @@ -342,31 +345,30 @@ func (a *RemoteUnloaderAdapter) subscribeProgress(nodeID, opID string, onProgres // Timeout: configured via DistributedConfig.BackendUpgradeTimeoutOrDefault // (default 15m). Real-world worst case observed: 8-10 minutes for large // CUDA-l4t backend images on Jetson over WiFi. -func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error) { - subject := messaging.SubjectNodeBackendUpgrade(nodeID) - xlog.Info("Sending NATS backend.upgrade", "nodeID", nodeID, "backend", backendType, "replica", replicaIndex, "opID", opID) - - sub := a.subscribeProgress(nodeID, opID, onProgress) - - reply, err := controlRequestJSON[workerctl.BackendUpgradeRequest, workerctl.BackendUpgradeReply](a.nats, subject, workerctl.BackendUpgradeRequest{ - Backend: backendType, - BackendGalleries: galleriesJSON, - URI: uri, - Name: name, - Alias: alias, - ReplicaIndex: int32(replicaIndex), - OpID: opID, - }, a.upgradeTimeout) - - if sub != nil { - if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { - xlog.Warn("Failed to unsubscribe from backend upgrade progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) - } - } - - if err != nil && isNATSTimeout(err) { - return nil, fmt.Errorf("%w (subject=%s nodeID=%s backend=%s): %v", - galleryop.ErrWorkerStillInstalling, subject, nodeID, backendType, err) +func (a *nodeControl) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error) { + target := a.link.describe(nodeID, workerctl.VerbBackendUpgrade) + xlog.Info("Sending "+a.link.label()+" backend.upgrade", "nodeID", nodeID, "backend", backendType, "replica", replicaIndex, "opID", opID) + + reply, err := callVerbWithProgress[workerctl.BackendUpgradeRequest, workerctl.BackendUpgradeReply]( + context.Background(), a.link, nodeID, workerctl.VerbBackendUpgrade, opID, workerctl.BackendUpgradeRequest{ + Backend: backendType, + BackendGalleries: galleriesJSON, + URI: uri, + Name: name, + Alias: alias, + ReplicaIndex: int32(replicaIndex), + OpID: opID, + }, a.upgradeTimeout, onProgress) + + if err != nil && errors.Is(err, errVerbNotServed) { + // A worker that answers, and has no such verb, is older than the verb. + // The caller falls back to the install it does understand, which is what + // ErrNoRoute has always meant for this verb. + return nil, fmt.Errorf("%w: %v", ErrNoRoute, err) + } + if err != nil && a.link.waitExpired(err) { + return nil, fmt.Errorf("%w (%s nodeID=%s backend=%s): %v", + galleryop.ErrWorkerStillInstalling, target, nodeID, backendType, err) } if err == nil { a.dropStoppedReplicaRows(nodeID, "backend.upgrade", backendType, reply.StoppedProcessKeys, reply.ReportsStoppedProcesses) @@ -380,42 +382,34 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO // doesn't subscribe to the new subject). It re-fires the legacy // backend.install with Force=true. Drop this once every worker is on // 2026-05-08 or newer. -func (a *RemoteUnloaderAdapter) InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { - subject := messaging.SubjectNodeBackendInstall(nodeID) +func (a *nodeControl) InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { + target := a.link.describe(nodeID, workerctl.VerbBackendInstall) xlog.Warn("Falling back to legacy backend.install Force=true (old worker)", "nodeID", nodeID, "backend", backendType) - sub := a.subscribeProgress(nodeID, opID, onProgress) - - reply, err := controlRequestJSON[workerctl.BackendInstallRequest, workerctl.BackendInstallReply](a.nats, subject, workerctl.BackendInstallRequest{ - Backend: backendType, - BackendGalleries: galleriesJSON, - URI: uri, - Name: name, - Alias: alias, - ReplicaIndex: int32(replicaIndex), - Force: true, - OpID: opID, - }, a.upgradeTimeout) - - if sub != nil { - if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { - xlog.Warn("Failed to unsubscribe from legacy backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) - } - } - - if err != nil && isNATSTimeout(err) { - return nil, fmt.Errorf("%w (subject=%s nodeID=%s backend=%s): %v", - galleryop.ErrWorkerStillInstalling, subject, nodeID, backendType, err) + reply, err := callVerbWithProgress[workerctl.BackendInstallRequest, workerctl.BackendInstallReply]( + context.Background(), a.link, nodeID, workerctl.VerbBackendInstall, opID, workerctl.BackendInstallRequest{ + Backend: backendType, + BackendGalleries: galleriesJSON, + URI: uri, + Name: name, + Alias: alias, + ReplicaIndex: int32(replicaIndex), + Force: true, + OpID: opID, + }, a.upgradeTimeout, onProgress) + + if err != nil && a.link.waitExpired(err) { + return nil, fmt.Errorf("%w (%s nodeID=%s backend=%s): %v", + galleryop.ErrWorkerStillInstalling, target, nodeID, backendType, err) } return reply, err } // ListBackends queries a worker node for its installed backends via NATS request-reply. -func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*workerctl.BackendListReply, error) { - subject := messaging.SubjectNodeBackendList(nodeID) - xlog.Debug("Sending NATS backend.list", "nodeID", nodeID) +func (a *nodeControl) ListBackends(nodeID string) (*workerctl.BackendListReply, error) { + xlog.Debug("Sending "+a.link.label()+" backend.list", "nodeID", nodeID) - return controlRequestJSON[workerctl.BackendListRequest, workerctl.BackendListReply](a.nats, subject, workerctl.BackendListRequest{}, 30*time.Second) + return callVerb[workerctl.BackendListRequest, workerctl.BackendListReply](context.Background(), a.link, nodeID, workerctl.VerbBackendList, workerctl.BackendListRequest{}, 30*time.Second) } // PingNode checks that a worker still has a live subscription on the bus. @@ -436,15 +430,12 @@ func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*workerctl.BackendL // A worker that answers anything is alive. Only when every subject reports no // route is the node treated as absent, so adding a newer subject here can // never condemn an older worker. -func (a *RemoteUnloaderAdapter) PingNode(nodeID string) error { - subjects := []string{ - messaging.SubjectNodeBackendList(nodeID), - messaging.SubjectNodeModelsRunning(nodeID), - } +func (a *nodeControl) PingNode(nodeID string) error { + verbs := []string{workerctl.VerbBackendList, workerctl.VerbModelsRunning} var lastErr error - for _, subject := range subjects { - _, err := controlRequestJSON[workerctl.BackendListRequest, workerctl.BackendListReply]( - a.nats, subject, workerctl.BackendListRequest{}, 5*time.Second) + for _, verb := range verbs { + _, err := callVerb[workerctl.BackendListRequest, workerctl.BackendListReply]( + context.Background(), a.link, nodeID, verb, workerctl.BackendListRequest{}, 5*time.Second) if err == nil { return nil } @@ -465,10 +456,9 @@ func (a *RemoteUnloaderAdapter) PingNode(nodeID string) error { // in-memory process table, so a slow reply means the worker itself is in // trouble, and the caller treats no-answer as "don't know" rather than as // "nothing running". -func (a *RemoteUnloaderAdapter) ListRunningModels(nodeID string) (*workerctl.ModelsRunningReply, error) { - subject := messaging.SubjectNodeModelsRunning(nodeID) - return controlRequestJSON[workerctl.ModelsRunningRequest, workerctl.ModelsRunningReply]( - a.nats, subject, workerctl.ModelsRunningRequest{}, 10*time.Second) +func (a *nodeControl) ListRunningModels(nodeID string) (*workerctl.ModelsRunningReply, error) { + return callVerb[workerctl.ModelsRunningRequest, workerctl.ModelsRunningReply]( + context.Background(), a.link, nodeID, workerctl.VerbModelsRunning, workerctl.ModelsRunningRequest{}, 10*time.Second) } // backendStopAckTimeout bounds the wait for a worker's backend.stop reply. @@ -483,7 +473,7 @@ const backendStopAckTimeout = 15 * time.Second // StopBackend tells a worker node to stop a specific gRPC backend process. // If backend is empty, the worker stops ALL backends. // The node stays registered and can receive another InstallBackend later. -func (a *RemoteUnloaderAdapter) StopBackend(nodeID, backend string) error { +func (a *nodeControl) StopBackend(nodeID, backend string) error { return a.stopBackend(nodeID, backend, false) } @@ -506,14 +496,13 @@ func (a *RemoteUnloaderAdapter) StopBackend(nodeID, backend string) error { // subscribes to this node's subject at all — is reported, because under the // old Publish it was reported too, and callers such as UnloadRemoteModel rely // on that to skip the registry cleanup for a node they could not reach. -func (a *RemoteUnloaderAdapter) stopBackend(nodeID, backend string, force bool) error { - subject := messaging.SubjectNodeBackendStop(nodeID) +func (a *nodeControl) stopBackend(nodeID, backend string, force bool) error { req := workerctl.BackendStopRequest{Backend: backend, Force: force} - reply, err := controlRequestJSON[workerctl.BackendStopRequest, workerctl.BackendStopReply]( - a.nats, subject, req, backendStopAckTimeout) + reply, err := callVerb[workerctl.BackendStopRequest, workerctl.BackendStopReply]( + context.Background(), a.link, nodeID, workerctl.VerbBackendStop, req, backendStopAckTimeout) if err != nil { - if isStrictRequestTimeout(err) { + if a.link.acknowledgementMissing(err) { xlog.Warn("Worker did not acknowledge backend.stop; assuming an older worker delivered it", "nodeID", nodeID, "backend", backend, "force", force) return nil @@ -539,11 +528,10 @@ func (a *RemoteUnloaderAdapter) stopBackend(nodeID, backend string, force bool) } // DeleteBackend tells a worker node to delete a backend (stop + remove files). -func (a *RemoteUnloaderAdapter) DeleteBackend(nodeID, backendName string) (*workerctl.BackendDeleteReply, error) { - subject := messaging.SubjectNodeBackendDelete(nodeID) - xlog.Info("Sending NATS backend.delete", "nodeID", nodeID, "backend", backendName) +func (a *nodeControl) DeleteBackend(nodeID, backendName string) (*workerctl.BackendDeleteReply, error) { + xlog.Info("Sending "+a.link.label()+" backend.delete", "nodeID", nodeID, "backend", backendName) - reply, err := controlRequestJSON[workerctl.BackendDeleteRequest, workerctl.BackendDeleteReply](a.nats, subject, workerctl.BackendDeleteRequest{Backend: backendName}, 2*time.Minute) + reply, err := callVerb[workerctl.BackendDeleteRequest, workerctl.BackendDeleteReply](context.Background(), a.link, nodeID, workerctl.VerbBackendDelete, workerctl.BackendDeleteRequest{Backend: backendName}, 2*time.Minute) if err != nil { return reply, err } @@ -567,7 +555,7 @@ func (a *RemoteUnloaderAdapter) DeleteBackend(nodeID, backendName string) (*work // then indistinguishable from "stopped nothing", so it must NOT be read as a // completed cleanup: leave the rows alone and fall back to the probe-based // staleness recovery that was the only mechanism before this change. -func (a *RemoteUnloaderAdapter) dropStoppedReplicaRows(nodeID, op, backendName string, processKeys []string, reported bool) { +func (a *nodeControl) dropStoppedReplicaRows(nodeID, op, backendName string, processKeys []string, reported bool) { if !reported { xlog.Debug("Worker did not report stopped processes; relying on probe-based staleness recovery", "nodeID", nodeID, "op", op, "backend", backendName) @@ -601,12 +589,12 @@ func (a *RemoteUnloaderAdapter) dropStoppedReplicaRows(nodeID, op, backendName s // // Callers that remove the replica row before they unload must pass the row they // read first. Looking the address up afterwards finds nothing. -func (a *RemoteUnloaderAdapter) UnloadReplica(nodeID string, replica NodeModel) error { +func (a *nodeControl) UnloadReplica(nodeID string, replica NodeModel) error { if replica.Address == "" { return nil } - xlog.Info("Sending NATS model.unload", "nodeID", nodeID, "model", replica.ModelName, "replica", replica.ReplicaIndex) - reply, err := controlRequestJSON[workerctl.ModelUnloadRequest, workerctl.ModelUnloadReply](a.nats, messaging.SubjectNodeModelUnload(nodeID), + xlog.Info("Sending "+a.link.label()+" model.unload", "nodeID", nodeID, "model", replica.ModelName, "replica", replica.ReplicaIndex) + reply, err := callVerb[workerctl.ModelUnloadRequest, workerctl.ModelUnloadReply](context.Background(), a.link, nodeID, workerctl.VerbModelUnload, workerctl.ModelUnloadRequest{ModelName: replica.ModelName, Address: replica.Address}, 30*time.Second) if err != nil { return err @@ -621,7 +609,7 @@ func (a *RemoteUnloaderAdapter) UnloadReplica(nodeID string, replica NodeModel) // each by address. It reads the replicas from the registry, so it only works // while their rows still exist. A caller that deletes the rows first uses // UnloadReplica with the row it read. -func (a *RemoteUnloaderAdapter) UnloadModelOnNode(nodeID, modelName string) error { +func (a *nodeControl) UnloadModelOnNode(nodeID, modelName string) error { lister, ok := a.registry.(interface { GetNodeModels(ctx context.Context, nodeID string) ([]NodeModel, error) }) @@ -645,7 +633,7 @@ func (a *RemoteUnloaderAdapter) UnloadModelOnNode(nodeID, modelName string) erro // DeleteModelFiles sends model.delete to all nodes that have the model cached. // This removes model files from worker disks. -func (a *RemoteUnloaderAdapter) DeleteModelFiles(modelName string) error { +func (a *nodeControl) DeleteModelFiles(modelName string) error { nodes, err := a.registry.FindNodesWithModel(context.Background(), modelName) if err != nil || len(nodes) == 0 { xlog.Debug("No nodes with model for file deletion", "model", modelName) @@ -653,10 +641,9 @@ func (a *RemoteUnloaderAdapter) DeleteModelFiles(modelName string) error { } for _, node := range nodes { - subject := messaging.SubjectNodeModelDelete(node.ID) - xlog.Info("Sending NATS model.delete", "nodeID", node.ID, "model", modelName) + xlog.Info("Sending "+a.link.label()+" model.delete", "nodeID", node.ID, "model", modelName) - reply, err := controlRequestJSON[workerctl.ModelDeleteRequest, workerctl.ModelDeleteReply](a.nats, subject, workerctl.ModelDeleteRequest{ModelName: modelName}, 30*time.Second) + reply, err := callVerb[workerctl.ModelDeleteRequest, workerctl.ModelDeleteReply](context.Background(), a.link, node.ID, workerctl.VerbModelDelete, workerctl.ModelDeleteRequest{ModelName: modelName}, 30*time.Second) if err != nil { xlog.Warn("model.delete failed on node", "node", node.Name, "error", err) continue @@ -673,3 +660,28 @@ func (a *RemoteUnloaderAdapter) StopNode(nodeID string) error { subject := messaging.SubjectNodeStop(nodeID) return a.nats.Publish(subject, nil) } + +// errVerbNotServed means that the worker answered and does not serve the verb. +// It is a fact about the version of the worker. It is not a routing fact and +// not an answer about any backend, so only a caller that has a fallback for an +// older worker reads it, and it becomes ErrNoRoute there. +var errVerbNotServed = errors.New("the worker does not serve that control verb") + +// callVerb sends one request and returns the decoded reply, or nil and the +// error. +func callVerb[Req, Reply any](ctx context.Context, link controlLink, nodeID, verb string, req Req, timeout time.Duration) (*Reply, error) { + var reply Reply + if err := link.request(ctx, nodeID, verb, req, &reply, timeout); err != nil { + return nil, err + } + return &reply, nil +} + +// callVerbWithProgress is callVerb for a verb that reports progress. +func callVerbWithProgress[Req, Reply any](ctx context.Context, link controlLink, nodeID, verb, opID string, req Req, timeout time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*Reply, error) { + var reply Reply + if err := link.requestWithProgress(ctx, nodeID, verb, opID, req, &reply, timeout, onProgress); err != nil { + return nil, err + } + return &reply, nil +} From bf045cf85df8e7c1d25a448f2ce1f3b15138dcdb Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:53:05 +0000 Subject: [PATCH 057/162] feat(cluster): publish the address and the peer credential of a replica A replica that holds the tunnel of a worker must be reachable by the other replicas, and they must be able to tell which replica dialled them. The instances row now holds the address that peers dial and the SHA-256 of a peer credential. A replica mints its own credential, publishes only the hash and keeps the secret, so the secret never crosses the network. The membership loop writes both in the same statement as the rest of the row, again after a sweep removed the row. A registration that carries neither leaves what an earlier one wrote. The address comes from the route to the database unless the operator sets one, and an address that cannot be dialled is refused with a reason (loopback and bind addresses are reported and not refused, because one host is a supported topology). Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/advertise.go | 189 +++++++++++++++++++++ core/services/cluster/advertise_test.go | 104 ++++++++++++ core/services/cluster/instance.go | 57 ++++++- core/services/cluster/membership.go | 20 ++- core/services/cluster/peeridentity.go | 71 ++++++++ core/services/cluster/peeridentity_test.go | 108 ++++++++++++ 6 files changed, 541 insertions(+), 8 deletions(-) create mode 100644 core/services/cluster/advertise.go create mode 100644 core/services/cluster/advertise_test.go create mode 100644 core/services/cluster/peeridentity.go create mode 100644 core/services/cluster/peeridentity_test.go diff --git a/core/services/cluster/advertise.go b/core/services/cluster/advertise.go new file mode 100644 index 000000000000..b3e66055bad9 --- /dev/null +++ b/core/services/cluster/advertise.go @@ -0,0 +1,189 @@ +package cluster + +import ( + "errors" + "fmt" + "net" + "net/url" + "strconv" + "strings" +) + +// DiscoverAdvertisedAddr determines the address this replica should advertise +// to its peers, with no operator configuration. +// +// Every replica in a deployment reaches the same PostgreSQL server, so the +// local interface that routes to PostgreSQL is on a network all the replicas +// demonstrably share. Opening a UDP socket toward the database sends no packet; +// it only asks the kernel to pick a source address for that route, which is the +// address to advertise. The caller supplies the port, since the frontend's +// listening port has nothing to do with the database's. +// +// What defeats the discovery is a DSN that NAMES loopback, not the database +// being co-located. Co-location is fine as long as the DSN names something +// routable: compose's usual `host=postgres` resolves to a bridge address, so +// the kernel picks this container's own bridge IP as the source, which is the +// address a peer on that network dials. It is `host=localhost` (or 127.0.0.1, +// or ::1) that makes the route loopback, and advertising 127.0.0.1 would make +// a peer dialling this replica reach itself instead. So an unspecified, +// loopback, or scoped source address is rejected with an error telling the +// operator to configure the advertised address explicitly, rather than +// returned. There is no fallback string: no address is better than a wrong one. +func DiscoverAdvertisedAddr(dsn string, port int) (string, error) { + // A port of 0 (or out of range) would produce an address nothing can dial, + // and the caller is likelier to have passed an unset field than to mean it. + if port < 1 || port > 65535 { + return "", fmt.Errorf("advertised port %d is out of range 1-65535", port) + } + host, dbPort, err := dsnHostPort(dsn) + if err != nil { + return "", err + } + conn, err := net.Dial("udp", net.JoinHostPort(host, dbPort)) + if err != nil { + return "", fmt.Errorf("resolving route to database host %q: %w", host, err) + } + // Nothing was ever sent on this socket, so a close failure carries no + // information about the address we just read. + defer func() { _ = conn.Close() }() + local, ok := conn.LocalAddr().(*net.UDPAddr) + if !ok || local.IP == nil { + return "", fmt.Errorf("no local address on the route to database host %q; set the advertised address explicitly", host) + } + if reason := unroutableReason(local.IP, local.Zone); reason != "" { + return "", fmt.Errorf("the route to database host %q is %s; set the advertised address explicitly", host, reason) + } + return net.JoinHostPort(local.IP.String(), strconv.Itoa(port)), nil +} + +// unroutableReason says why ip cannot serve as an address other hosts dial, or +// "" when it can. It is the one place that decides, so the discovered address +// and the configured one are held to the same rule; they differ only in what +// they do with the answer. +func unroutableReason(ip net.IP, zone string) string { + switch { + case ip == nil || ip.IsUnspecified(): + return fmt.Sprintf("unspecified (%s), which is a bind address rather than one anything can connect to", ip) + case ip.IsLoopback(): + return fmt.Sprintf("loopback (%s), which means \"this host\" to whoever dials it, so every peer would reach itself", ip) + case ip.IsLinkLocalUnicast(): + return fmt.Sprintf("link-local (%s), which peers on other hosts cannot dial", withZone(ip, zone)) + // A zone is normally attached only to a link-local address, which the case + // above already rejects. This one stays for the scoped address of some + // other class a platform may hand back, and says so rather than repeating + // the link-local label: the two have different cures, and an operator told + // the wrong one looks in the wrong place. + case zone != "": + return fmt.Sprintf("scoped to interface %q (%s), and the zone is dropped by the time an address is stored, leaving a host nothing can dial", zone, withZone(ip, zone)) + } + return "" +} + +// withZone renders the address the way it has to be dialled. IP.String() drops +// the %iface, so an unadorned %s in a rejection reports an address that differs +// from the one being rejected. +func withZone(ip net.IP, zone string) string { + if zone == "" { + return ip.String() + } + return ip.String() + "%" + zone +} + +// CheckAdvertisedAddr validates an address an operator configured, returning a +// reason it is questionable, or an error if it is unusable. +// +// A configured address bypasses every check DiscoverAdvertisedAddr performs, +// and the value most likely to be copied is the one that works on a single +// host: "127.0.0.1:8080" on three hosts makes every peer dial itself, which +// presents as a relay loop rather than as a configuration error. +// +// The split between error and reason is deliberate. An address that cannot be +// parsed into host and port is an error, because nothing can dial it at all. An +// address that merely means "this host" is a reason to warn and no more: a +// single-host deployment, including this repository's own e2e cluster, uses one +// correctly, and refusing it would be refusing a supported topology. +func CheckAdvertisedAddr(addr string) (reason string, err error) { + host, port, err := net.SplitHostPort(addr) + if err != nil { + return "", fmt.Errorf("advertised address %q is not host:port: %w", addr, err) + } + if host == "" { + return "", fmt.Errorf("advertised address %q names no host, so peers have nothing to dial", addr) + } + portNumber, err := strconv.Atoi(port) + if err != nil || portNumber < 1 || portNumber > 65535 { + return "", fmt.Errorf("advertised address %q has no usable port (want 1-65535)", addr) + } + // The zone is split off before parsing because net.ParseIP rejects + // "fe80::1%eth0" outright. Left joined, a scoped literal would look like a + // name and collect no warning at all, which is the one case where the + // address is guaranteed not to work for a peer. + host, zone := splitZone(host) + // A name is resolved by whoever dials it, and may resolve differently + // there, so its presence is all this side can check. + ip := net.ParseIP(host) + if ip == nil { + return "", nil + } + return unroutableReason(ip, zone), nil +} + +// splitZone separates an IPv6 scope from the address it qualifies. A name +// never carries one, so a host with no "%" comes back unchanged. +func splitZone(host string) (string, string) { + addr, zone, found := strings.Cut(host, "%") + if !found { + return host, "" + } + return addr, zone +} + +// dsnHostPort extracts the host and port from either DSN form gorm's postgres +// driver accepts: a URL ("postgres://user:pass@host:5432/db") or libpq keyword +// pairs ("host=... port=..."). +func dsnHostPort(dsn string) (string, string, error) { + const defaultPort = "5432" + dsn = strings.TrimSpace(dsn) + if dsn == "" { + return "", "", errors.New("empty database DSN") + } + + if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") { + u, err := url.Parse(dsn) + if err != nil { + return "", "", fmt.Errorf("parsing database DSN: %w", err) + } + host := u.Hostname() + if host == "" { + return "", "", errors.New("database DSN has no host") + } + port := u.Port() + if port == "" { + port = defaultPort + } + return host, port, nil + } + + host, port := "", defaultPort + for _, field := range strings.Fields(dsn) { + key, value, found := strings.Cut(field, "=") + if !found { + continue + } + switch key { + case "host": + host = value + case "port": + port = value + } + } + if host == "" { + return "", "", errors.New("database DSN has no host") + } + // A Unix socket directory tells us nothing about which interface reaches + // the database, so there is no address to derive. + if strings.HasPrefix(host, "/") { + return "", "", fmt.Errorf("database DSN uses a unix socket (%q); no routable address to advertise", host) + } + return host, port, nil +} diff --git a/core/services/cluster/advertise_test.go b/core/services/cluster/advertise_test.go new file mode 100644 index 000000000000..92a31815dff0 --- /dev/null +++ b/core/services/cluster/advertise_test.go @@ -0,0 +1,104 @@ +package cluster_test + +import ( + "net" + + "github.com/mudler/LocalAI/core/services/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Advertised address discovery", func() { + // The address itself depends on host networking and is deliberately not + // asserted. What is portable is the shape: whatever interface routes to the + // database, the port must be the one the caller asked for, not the + // database's. + It("combines a local interface with the caller's port", func() { + addr, err := cluster.DiscoverAdvertisedAddr("postgres://198.51.100.1:5432/testdb", 8080) + if err != nil { + Skip("no route to a database host on this machine: " + err.Error()) + } + host, port, splitErr := net.SplitHostPort(addr) + Expect(splitErr).ToNot(HaveOccurred()) + Expect(port).To(Equal("8080")) + Expect(net.ParseIP(host)).ToNot(BeNil()) + }) + + It("refuses a DSN it cannot derive an address from", func() { + _, err := cluster.DiscoverAdvertisedAddr("", 8080) + Expect(err).To(HaveOccurred()) + }) + + // A DSN that NAMES loopback routes over loopback on every platform, so this + // is deterministic rather than host-dependent. Co-location is not the + // trigger: compose's `host=postgres` resolves to a bridge address and + // discovery works there. Returning 127.0.0.1 would make a peer dialling + // this replica reach itself. + It("refuses a loopback route instead of advertising an address peers cannot use", func() { + addr, err := cluster.DiscoverAdvertisedAddr("postgres://user@127.0.0.1:5432/testdb", 8080) + Expect(addr).To(BeEmpty()) + Expect(err).To(MatchError(ContainSubstring("loopback"))) + }) + + It("refuses a port that cannot be dialled", func() { + _, err := cluster.DiscoverAdvertisedAddr("postgres://198.51.100.1:5432/testdb", 0) + Expect(err).To(MatchError(ContainSubstring("out of range"))) + }) +}) + +var _ = Describe("Checking a configured advertised address", func() { + // The configured address bypasses discovery entirely, so it bypasses every + // rejection discovery makes. These are the checks that put back the ones + // that can be made without a route to look at. + It("accepts an address on a network other hosts can reach", func() { + reason, err := cluster.CheckAdvertisedAddr("10.0.0.7:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(reason).To(BeEmpty()) + }) + + It("accepts a name, because the dialler is what resolves it", func() { + reason, err := cluster.CheckAdvertisedAddr("localai-frontend.default.svc:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(reason).To(BeEmpty()) + }) + + It("refuses an address with no port, which nothing could dial", func() { + _, err := cluster.CheckAdvertisedAddr("10.0.0.7") + Expect(err).To(HaveOccurred()) + }) + + It("refuses a port outside the dialable range", func() { + _, err := cluster.CheckAdvertisedAddr("10.0.0.7:0") + Expect(err).To(MatchError(ContainSubstring("port"))) + }) + + It("refuses an address that names no host", func() { + _, err := cluster.CheckAdvertisedAddr(":8080") + Expect(err).To(MatchError(ContainSubstring("no host"))) + }) + + It("reports loopback without refusing it, because one host is a supported topology", func() { + // Correct on a single host, and the value most likely to be copied + // onto three, where every peer would then dial itself. + reason, err := cluster.CheckAdvertisedAddr("127.0.0.1:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(reason).To(ContainSubstring("loopback")) + }) + + It("reports a bind address, which is not an address at all", func() { + reason, err := cluster.CheckAdvertisedAddr("0.0.0.0:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(reason).To(ContainSubstring("unspecified")) + }) + + It("reports a scoped literal, which net.ParseIP alone would wave through as a name", func() { + // The zone has to be split off before parsing, or this address is + // indistinguishable from a hostname and collects no warning at all. + reason, err := cluster.CheckAdvertisedAddr("[fe80::1%eth0]:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(reason).ToNot(BeEmpty(), "a scoped address peers cannot dial was accepted in silence") + Expect(reason).To(ContainSubstring("fe80::1%eth0"), + "the reported address must carry its zone, or it is not the address being rejected") + }) +}) diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index df2290aae6b3..47286fdce1eb 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -29,6 +29,15 @@ type Instance struct { // ReadyReason says why the replica is not ready. It is empty when the // replica is ready. ReadyReason string `gorm:"size:512;not null;default:''" json:"ready_reason,omitempty"` + // AdvertisedAddr is the host and port at which the other replicas dial this + // one. It is empty for a replica that publishes none, which no peer can + // reach. + AdvertisedAddr string `gorm:"size:255" json:"advertised_addr,omitempty"` + // PeerTokenHash is the SHA-256 of the peer credential of this replica, as + // hex. A peer that dials this replica proves who it is by presenting the + // secret behind this hash. It is empty for a replica that has not published + // one, and an empty hash authorises nobody. + PeerTokenHash string `gorm:"size:64" json:"-"` } // Registry reads and writes the instances table and the table of worker @@ -48,6 +57,17 @@ func NewRegistry(db *gorm.DB) *Registry { // live set never sees a live replica as missing. The readiness columns are written too, so a replica // that registers again after it was swept does not lose what it reported. func (r *Registry) Register(ctx context.Context, id, version string, readyEpoch int64, readyReason string) error { + return r.RegisterPeer(ctx, id, version, readyEpoch, readyReason, "", "") +} + +// RegisterPeer is Register for a replica that other replicas dial. The address +// and the hash of the peer credential go in the same statement as the rest of the +// row. A replica that had an address and no identity, even for a moment, would be +// one that every peer refuses. +// +// An empty address or hash leaves the stored value alone on an update, so a +// replica that publishes none cannot wipe what an earlier registration wrote. +func (r *Registry) RegisterPeer(ctx context.Context, id, version string, readyEpoch int64, readyReason, advertisedAddr, peerTokenHash string) error { // The database stamps last_seen and not this process. Liveness is compared // across replicas, so it must be measured on the one clock they share. With // one clock per replica, the real window becomes the window minus the clock @@ -59,20 +79,43 @@ func (r *Registry) Register(ctx context.Context, id, version string, readyEpoch "ready_reason": readyReason, "last_seen": gorm.Expr("now()"), } + update := map[string]any{ + "version": version, + "ready_epoch": readyEpoch, + "ready_reason": readyReason, + "last_seen": gorm.Expr("now()"), + } + if advertisedAddr != "" { + row["advertised_addr"] = advertisedAddr + update["advertised_addr"] = advertisedAddr + } + if peerTokenHash != "" { + row["peer_token_hash"] = peerTokenHash + update["peer_token_hash"] = peerTokenHash + } if err := r.db.WithContext(ctx).Model(&Instance{}).Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "id"}}, - DoUpdates: clause.Assignments(map[string]any{ - "version": version, - "ready_epoch": readyEpoch, - "ready_reason": readyReason, - "last_seen": gorm.Expr("now()"), - }), + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.Assignments(update), }).Create(row).Error; err != nil { return fmt.Errorf("registering instance %q: %w", id, err) } return nil } +// Get returns one instance, or ErrInstanceNotFound if it is not registered. It +// does not ask if the replica is alive. +func (r *Registry) Get(ctx context.Context, id string) (*Instance, error) { + var inst Instance + err := r.db.WithContext(ctx).Where("id = ?", id).First(&inst).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, fmt.Errorf("getting instance %q: %w", id, ErrInstanceNotFound) + } + if err != nil { + return nil, fmt.Errorf("getting instance %q: %w", id, err) + } + return &inst, nil +} + // Heartbeat refreshes LastSeen for a registered instance. An unknown ID is an // error and not an insert, because the caller must then register again with // the full row. diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index 7e72e901cbf7..90920490e9dd 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -103,6 +103,8 @@ type Membership struct { readyEpoch int64 readyReason string reclaimer Reclaimer + advertised string + peerHash string } // NewMembership returns the membership loop for one replica. @@ -134,6 +136,21 @@ func (m *Membership) SetReconnectGrace(grace time.Duration) { m.retention = DepartedRetentionFor(grace) } +// SetPeer gives the loop what the other replicas need to dial this one: the +// address that they dial, and the credential that proves who this replica is. +// Every registration writes both, so a replica that a peer swept and that +// registers again is dialled again. Set it before Start. +// +// It is safe on a nil receiver, like Stop. +func (m *Membership) SetPeer(advertisedAddr string, cred PeerCredential) { + if m == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + m.advertised, m.peerHash = advertisedAddr, cred.Hash() +} + // SetReclaimer gives the loop the holder of the worker connections of this // replica. After a sweep deleted the row of this replica, the loop registers it // again and asks the holder to claim its connections again. Without this, the @@ -169,8 +186,9 @@ func (m *Membership) ReportReady(ctx context.Context, epoch int64, reason string func (m *Membership) register(ctx context.Context) error { m.mu.Lock() epoch, reason := m.readyEpoch, m.readyReason + advertised, peerHash := m.advertised, m.peerHash m.mu.Unlock() - return m.reg.Register(ctx, m.id, m.version, epoch, reason) + return m.reg.RegisterPeer(ctx, m.id, m.version, epoch, reason, advertised, peerHash) } // Start registers this replica, then heartbeats and sweeps in the background. diff --git a/core/services/cluster/peeridentity.go b/core/services/cluster/peeridentity.go new file mode 100644 index 000000000000..e6c8319e116a --- /dev/null +++ b/core/services/cluster/peeridentity.go @@ -0,0 +1,71 @@ +package cluster + +import ( + "crypto/rand" + "crypto/sha256" + "crypto/subtle" + "encoding/hex" +) + +// PeerIdentityHeader carries the peer credential of a replica when it dials +// another replica at GET /api/cluster/peer. +// +// The credential of a replica is the only secret that opens a peer link. The +// registration token of the deployment is for workers that join the cluster, and +// it opens no peer link: a replica that is empty of that token has the same +// peers as one that has it, and a worker that holds it cannot pose as a replica. +const PeerIdentityHeader = "X-LocalAI-Peer-Token" + +// PeerCredential is the proof of one frontend replica of which replica it is. +// +// It follows the per-node credential of a worker: a random secret, stored only as +// its SHA-256, compared in constant time, with no fallback to a shared token. It +// differs in one way, in the stronger direction. A replica writes its own row in +// the instances table, so it mints its secret itself, publishes only the hash, +// and the secret never leaves the process that made it. +// +// It is one value with both halves because they must not be minted twice. The +// membership loop stores Hash and the peer pool presents Token. Two mints would +// leave a replica whose stored hash and presented secret disagree, and every dial +// it makes would be refused. +// +// The zero value is a credential that nobody can use. An empty Token is refused +// by every handler and an empty Hash authorises nobody. +type PeerCredential struct { + token string + hash string +} + +// NewPeerCredential mints a credential for this replica. crypto/rand.Text gives +// at least 128 bits of randomness. +func NewPeerCredential() PeerCredential { + token := rand.Text() + return PeerCredential{token: token, hash: HashPeerToken(token)} +} + +// Token is the secret that this replica presents when it dials a peer. It is sent +// and never stored. +func (c PeerCredential) Token() string { return c.token } + +// Hash is what this replica publishes in its instances row, and the only form of +// the credential that another replica sees. +func (c PeerCredential) Hash() string { return c.hash } + +// HashPeerToken renders a presented token in the form that the instances table +// stores. The handler that checks a dial and the replica that registers its row +// call it, so they cannot disagree about the encoding. +func HashPeerToken(token string) string { + sum := sha256.Sum256([]byte(token)) + return hex.EncodeToString(sum[:]) +} + +// PeerTokenMatches reports whether a token is the one behind storedHash, in +// constant time. An empty hash matches nothing, and so does an empty token: a +// replica with no published credential authorises nobody, and "no credential" +// is never read as "any credential". +func PeerTokenMatches(token, storedHash string) bool { + if storedHash == "" || token == "" { + return false + } + return subtle.ConstantTimeCompare([]byte(HashPeerToken(token)), []byte(storedHash)) == 1 +} diff --git a/core/services/cluster/peeridentity_test.go b/core/services/cluster/peeridentity_test.go new file mode 100644 index 000000000000..ecdda67e76aa --- /dev/null +++ b/core/services/cluster/peeridentity_test.go @@ -0,0 +1,108 @@ +package cluster_test + +import ( + "context" + "strings" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Peer credential", func() { + It("mints a secret of at least 128 bits and publishes only its hash", func() { + cred := cluster.NewPeerCredential() + Expect(len(cred.Token())).To(BeNumerically(">=", 26), "crypto/rand.Text gives 26 characters of base32") + Expect(cred.Hash()).To(HaveLen(64)) + Expect(cred.Hash()).ToNot(ContainSubstring(cred.Token())) + Expect(cluster.HashPeerToken(cred.Token())).To(Equal(cred.Hash())) + }) + + It("mints a different secret each time", func() { + Expect(cluster.NewPeerCredential().Token()).ToNot(Equal(cluster.NewPeerCredential().Token())) + }) + + It("matches the secret behind a hash and no other", func() { + cred := cluster.NewPeerCredential() + Expect(cluster.PeerTokenMatches(cred.Token(), cred.Hash())).To(BeTrue()) + Expect(cluster.PeerTokenMatches(cluster.NewPeerCredential().Token(), cred.Hash())).To(BeFalse()) + Expect(cluster.PeerTokenMatches(strings.ToLower(cred.Token()), cred.Hash())).To(BeFalse()) + }) + + It("authorises nobody when the hash is empty, or when the token is", func() { + Expect(cluster.PeerTokenMatches("", "")).To(BeFalse(), "an empty credential is not a credential that matches everything") + Expect(cluster.PeerTokenMatches("anything", "")).To(BeFalse()) + Expect(cluster.PeerTokenMatches("", cluster.NewPeerCredential().Hash())).To(BeFalse()) + }) + + It("is unusable in its zero value", func() { + var zero cluster.PeerCredential + Expect(zero.Token()).To(BeEmpty()) + Expect(zero.Hash()).To(BeEmpty()) + Expect(cluster.PeerTokenMatches(zero.Token(), zero.Hash())).To(BeFalse()) + }) +}) + +var _ = Describe("Registering a replica that peers dial", func() { + var ( + ctx context.Context + reg *cluster.Registry + ) + + BeforeEach(func() { + ctx = context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("stores the address and the hash with the rest of the row, and reads them back", func() { + cred := cluster.NewPeerCredential() + Expect(reg.RegisterPeer(ctx, "a", "v1", 0, "", "10.0.0.7:8080", cred.Hash())).To(Succeed()) + got, err := reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + Expect(got.AdvertisedAddr).To(Equal("10.0.0.7:8080")) + Expect(got.PeerTokenHash).To(Equal(cred.Hash())) + }) + + It("replaces both when the replica registers again with new ones", func() { + Expect(reg.RegisterPeer(ctx, "a", "v1", 0, "", "10.0.0.7:8080", cluster.NewPeerCredential().Hash())).To(Succeed()) + again := cluster.NewPeerCredential() + Expect(reg.RegisterPeer(ctx, "a", "v1", 0, "", "10.0.0.8:8080", again.Hash())).To(Succeed()) + got, err := reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + Expect(got.AdvertisedAddr).To(Equal("10.0.0.8:8080")) + Expect(got.PeerTokenHash).To(Equal(again.Hash())) + }) + + It("does not wipe the address or the hash when a registration carries none", func() { + cred := cluster.NewPeerCredential() + Expect(reg.RegisterPeer(ctx, "a", "v1", 0, "", "10.0.0.7:8080", cred.Hash())).To(Succeed()) + Expect(reg.Register(ctx, "a", "v2", 3, "reason")).To(Succeed()) + got, err := reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + Expect(got.Version).To(Equal("v2")) + Expect(got.AdvertisedAddr).To(Equal("10.0.0.7:8080")) + Expect(got.PeerTokenHash).To(Equal(cred.Hash())) + }) + + It("writes them again when the membership loop registers a swept replica again", func() { + cred := cluster.NewPeerCredential() + m := cluster.NewMembership(reg, "swept", "v1") + m.SetPeer("10.0.0.9:8080", cred) + Expect(m.Start(ctx)).To(Succeed()) + DeferCleanup(m.Stop) + + got, err := reg.Get(ctx, "swept") + Expect(err).ToNot(HaveOccurred()) + Expect(got.AdvertisedAddr).To(Equal("10.0.0.9:8080")) + Expect(got.PeerTokenHash).To(Equal(cred.Hash())) + }) + + It("reports a replica that is not registered as such, and never as unreachable", func() { + _, err := reg.Get(ctx, "ghost") + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + }) +}) From 2b63ab59b84864435d9fd7c468c73e4b4d1eccd6 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:53:05 +0000 Subject: [PATCH 058/162] feat(tunnel): dial a worker through the replica that holds its tunnel WorkerDialer is the one door to a worker. A replica that holds the tunnel opens the stream itself. Any other replica relays through the owner over a link between replicas, so a worker that lands on one replica is reached through all of them. The failures stay apart. A route that does not exist is ErrNoRoute, with absence claims (no live owner, no such replica) kept out of the chain so that no caller can read them as a missing worker. A worker that refuses a stream keeps its refusal and carries no ErrNoRoute, because it has shown that it is there. A peer that does not answer is ErrPeerUnreachable and not ErrNoRoute: a link between two frontends is not a fact about a worker, and a scheduler that demoted the worker for it would lose capacity for nothing. An expired budget of the caller is checked first and is never a route failure. A dial can ask for the bulk lane, and then it never falls back to the lane of model calls. The request frame of the relay carries the lane, and the remaining time of the caller so that the owner bounds its own open. The dialling side of a link accepts no stream. The accepting side refuses a stream that does not name a worker it holds, and never relays onward. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/dialer.go | 246 +++++++++++++++ core/services/tunnel/dialer_test.go | 440 +++++++++++++++++++++++++++ core/services/tunnel/export_test.go | 5 + core/services/tunnel/peerlink.go | 320 +++++++++++++++++++ core/services/tunnel/peersessions.go | 141 +++++++++ core/services/tunnel/relay.go | 361 ++++++++++++++++++++++ core/services/tunnel/relay_test.go | 265 ++++++++++++++++ core/services/tunnel/route.go | 87 ++++++ core/services/tunnel/session.go | 20 +- 9 files changed, 1883 insertions(+), 2 deletions(-) create mode 100644 core/services/tunnel/dialer.go create mode 100644 core/services/tunnel/dialer_test.go create mode 100644 core/services/tunnel/peerlink.go create mode 100644 core/services/tunnel/peersessions.go create mode 100644 core/services/tunnel/relay.go create mode 100644 core/services/tunnel/relay_test.go create mode 100644 core/services/tunnel/route.go diff --git a/core/services/tunnel/dialer.go b/core/services/tunnel/dialer.go new file mode 100644 index 000000000000..2df40ff23acd --- /dev/null +++ b/core/services/tunnel/dialer.go @@ -0,0 +1,246 @@ +package tunnel + +import ( + "context" + "errors" + "fmt" + "net" + "time" + + "github.com/mudler/xlog" +) + +// PeerOpener opens a stream to another frontend replica. *PeerPool is the +// production implementation. It is an interface so that a spec can decide what a +// peer does without a second frontend. +type PeerOpener interface { + Open(ctx context.Context, peerID string) (net.Conn, error) +} + +// dialHandshakeTimeout bounds the request and reply that open every stream, when +// the caller stated no deadline of its own. It is a backstop and not a budget. A +// worker or a relay that accepted a stream and then said nothing would otherwise +// park the caller until the keepalive of the session killed it. A deadline of the +// caller is used instead when it is the shorter of the two. +const dialHandshakeTimeout = 15 * time.Second + +// WorkerDialer opens connections to a worker through its tunnel, wherever in the +// deployment that tunnel is held. +// +// It is the single door. A worker holds one tunnel, it lands on one replica, and +// no other part of the frontend dials the address of a worker. Every protocol +// that the frontend speaks to a worker goes through here: gRPC to a backend +// process, HTTP for control and file transfer, and a WebSocket for logs. A worker +// behind NAT has no address to dial, and a direct dial that works in a test with +// one replica fails in production. +type WorkerDialer struct { + // The registry carries both the identity of this replica and the sessions it + // holds, so that the identity this dialer compares owners against is the one + // the registry claims as. Two identities would make this replica relay to + // itself for every worker it holds. + tunnels *Registry + peers PeerOpener +} + +// NewWorkerDialer returns the dialer for the tunnels this replica holds and the +// peer links it can relay over. A nil peers means that this replica cannot +// relay, which is reported as ErrNoRelayPath and not as a worker being absent. +func NewWorkerDialer(tunnels *Registry, peers PeerOpener) *WorkerDialer { + return &WorkerDialer{tunnels: tunnels, peers: peers} +} + +// DialOption changes how Dial picks the session of the worker. +type DialOption func(*dialOptions) + +type dialOptions struct { + lane Lane + noFallback bool +} + +// WithBulkLane sends the stream over the bulk lane, and fails when the worker has +// no bulk session. It does not fall back to the inference lane. A file transfer +// asks for it: a transfer on the lane of model calls delays every one of them, +// which is what the bulk lane is for. +func WithBulkLane() DialOption { + return func(o *dialOptions) { o.lane, o.noFallback = LaneBulk, true } +} + +// Dial opens one stream to a local service of a worker. tag says which service +// (StreamTagGRPC or StreamTagHTTP) and target which instance of it. +// +// The connection that comes back is past both handshakes and carries the +// protocol of the caller and nothing else. No deadline is armed on it: what +// follows can be an inference that is quiet for minutes. +// +// Every failure to resolve or open a route carries ErrNoRoute, and no failure +// carries an absence sentinel. A worker that answers keeps its own refusal and no +// ErrNoRoute, because it has shown that it is there. See routeFailure for what is +// neither. +func (d *WorkerDialer) Dial(ctx context.Context, nodeID, tag, target string, opts ...DialOption) (net.Conn, error) { + o := dialOptions{lane: LaneInference} + for _, opt := range opts { + opt(&o) + } + var openOpts []OpenOption + if o.noFallback { + openOpts = append(openOpts, WithoutFallback()) + } + stream, err := d.tunnels.Open(ctx, nodeID, o.lane, openOpts...) + if err == nil { + return d.handshake(ctx, stream, nodeID, tag, target) + } + if !errors.Is(err, ErrNotOwner) { + // The tunnel is held here and would not carry a stream. ErrNotOwner stays + // out: it would send the caller to resolve an owner that is this replica. + return nil, routeFailure(nodeID, err) + } + return d.relay(ctx, nodeID, tag, target, o) +} + +// DialerFor returns a function in the shape of http.Transport.DialContext and +// websocket.Dialer.NetDialContext, bound to one worker and one local service of +// it. The network is ignored and the address becomes the target of the stream, +// which is what lets an http.Client reach the worker without knowing that a +// tunnel exists. What the worker does with the target is its own decision: the +// grpc tag keeps the port and dials its own loopback, and the http tag ignores +// the target. A frontend therefore cannot steer a dial of the worker. +func (d *WorkerDialer) DialerFor(nodeID, tag string, opts ...DialOption) func(ctx context.Context, network, addr string) (net.Conn, error) { + return func(ctx context.Context, _, addr string) (net.Conn, error) { + return d.Dial(ctx, nodeID, tag, addr, opts...) + } +} + +// GRPCDialerFor returns a function in the shape of grpc.WithContextDialer, bound +// to the backend processes of one worker. +func (d *WorkerDialer) GRPCDialerFor(nodeID string) func(ctx context.Context, addr string) (net.Conn, error) { + return func(ctx context.Context, addr string) (net.Conn, error) { + return d.Dial(ctx, nodeID, StreamTagGRPC, addr) + } +} + +// relay opens the stream through the replica that holds the tunnel of the worker. +func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o dialOptions) (net.Conn, error) { + // Owner and not the row: a row outlives its owner by up to a liveness window, + // and a dial to what the unjoined read returns would reach a process that is + // gone and report the worker as unreachable instead of absent. + owner, _, err := d.tunnels.reg.Owner(ctx, nodeID) + if err != nil { + // ErrNoConnection is the ordinary answer for a worker that has not dialled + // in yet, and it must not get out: routeFailure keeps it in the message. + return nil, routeFailure(nodeID, err) + } + if owner == d.tunnels.selfID { + // The table names this replica, and the registry says that the tunnel is + // not held here, so the attachment went away between the claim and now. + // Relaying would send the request into this same process. + return nil, routeFailure(nodeID, fmt.Errorf("the connection row names this replica, which no longer holds the tunnel: %w", ErrNotOwner)) + } + if d.peers == nil { + return nil, routeFailure(nodeID, fmt.Errorf("the tunnel is held by replica %q: %w", owner, ErrNoRelayPath)) + } + + stream, err := d.peers.Open(ctx, owner) + if err != nil { + // ErrPeerUnreachable and ErrPoolClosed keep their identity. An owner swept + // between the lookup and this dial is an absence claim about a replica, and + // routeFailure withholds it. + return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) + } + + // The remaining time of the caller, so that the owner can bound its own open. + if err := WriteRelayRequest(stream, nodeID, relayLane(o), remainingBudget(ctx)); err != nil { + _ = stream.Close() + return nil, routeFailure(nodeID, fmt.Errorf("naming the node on a stream to replica %q: %w", owner, err)) + } + if err := ReadRelayReply(stream); err != nil { + // A refusal of the owning replica and not of the worker: the owner would + // not relay, which is a route that does not exist. + _ = stream.Close() + if blamed := callerRanOut(ctx); blamed != nil { + return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, blamed)) + } + return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) + } + return d.handshake(ctx, stream, nodeID, tag, target) +} + +// handshake names the service of the worker on a stream and waits for the answer +// of the worker. It owns closing the stream on every failure: a stream left +// open after a failed handshake holds a slot of the session for as long as the +// session lives. +func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, tag, target string) (net.Conn, error) { + // The budget of the caller is checked first, before anything reads as the + // answer of a worker. The deadline of the handshake is the deadline of the + // caller when that is shorter, so a caller that runs out makes the timer of the + // socket fire, and the timeout arrives here while ctx.Err may still be nil. + // A refusal that arrives in the same instant is the timeout of the caller, and + // never evidence about a backend. + blameCaller := func(what string) error { + ctxErr := callerRanOut(ctx) + if ctxErr == nil { + return nil + } + return routeFailure(nodeID, fmt.Errorf("%s: the budget of the caller ran out: %w", what, ctxErr)) + } + + if err := stream.SetDeadline(handshakeDeadline(ctx)); err != nil { + _ = stream.Close() + return nil, routeFailure(nodeID, fmt.Errorf("arming the handshake deadline: %w", err)) + } + + if err := WriteStreamRequest(stream, tag, target); err != nil { + // The tunnel broke under the request. Nothing was asked of the worker and + // nothing was learned about it. + _ = stream.Close() + if blamed := blameCaller(fmt.Sprintf("asking for %q on %q", tag, target)); blamed != nil { + return nil, blamed + } + return nil, routeFailure(nodeID, fmt.Errorf("asking for %q on %q: %w", tag, target, err)) + } + if err := ReadStreamReply(stream); err != nil { + _ = stream.Close() + if blamed := blameCaller(fmt.Sprintf("opening %q", tag)); blamed != nil { + return nil, blamed + } + if IsWorkerAnswer(err) { + // The worker wrote a refusal, so it is connected and answering. This is + // the one failure on the path that is evidence about the worker, and the + // umbrella would throw it away. + return nil, fmt.Errorf("opening %q on node %q: %w", tag, nodeID, err) + } + return nil, routeFailure(nodeID, fmt.Errorf("opening %q: %w", tag, err)) + } + + // Cleared whatever the caller's context carried. What follows is the protocol + // of the caller, and a deadline of the handshake would end it. + if err := stream.SetDeadline(time.Time{}); err != nil { + _ = stream.Close() + return nil, routeFailure(nodeID, fmt.Errorf("clearing the handshake deadline: %w", err)) + } + xlog.Debug("opened a tunnelled stream to a worker", "node", nodeID, "tag", tag, "target", target) + return stream, nil +} + +// handshakeDeadline is when the handshake must be done: the deadline of the caller +// when it has one and it is sooner, and the backstop otherwise. +func handshakeDeadline(ctx context.Context) time.Time { + backstop := time.Now().Add(dialHandshakeTimeout) + deadline, ok := ctx.Deadline() + if !ok || deadline.After(backstop) { + return backstop + } + return deadline +} + +// remainingBudget is how long the caller is still willing to wait, or zero when +// it did not say. It is zero and not negative for an expired context. +func remainingBudget(ctx context.Context) time.Duration { + deadline, ok := ctx.Deadline() + if !ok { + return 0 + } + if remaining := time.Until(deadline); remaining > 0 { + return remaining + } + return 0 +} diff --git a/core/services/tunnel/dialer_test.go b/core/services/tunnel/dialer_test.go new file mode 100644 index 000000000000..d54cf301ab4b --- /dev/null +++ b/core/services/tunnel/dialer_test.go @@ -0,0 +1,440 @@ +package tunnel_test + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// scriptedWorker answers the request frame of every stream of a session with a +// script, and then echoes what it reads. It stands for the worker side of a +// tunnel; the session under it is a real websocket with a real yamux session. +type scriptedWorker struct { + mu sync.Mutex + requests []workerRequest + refuse error + hold time.Duration +} + +type workerRequest struct{ tag, target string } + +func (w *scriptedWorker) serve(sess *tunnel.Session) { + go func() { + for { + st, err := sess.AcceptStream() + if err != nil { + return + } + go func() { + defer func() { _ = st.Close() }() + tag, target, err := tunnel.ReadStreamRequest(st) + if err != nil { + return + } + w.mu.Lock() + w.requests = append(w.requests, workerRequest{tag, target}) + refuse, hold := w.refuse, w.hold + w.mu.Unlock() + if hold > 0 { + time.Sleep(hold) + } + if refuse != nil { + _ = tunnel.WriteStreamRefusal(st, refuse) + return + } + if err := tunnel.WriteStreamAccepted(st); err != nil { + return + } + _, _ = io.Copy(st, st) + }() + } + }() +} + +func (w *scriptedWorker) seen() []workerRequest { + w.mu.Lock() + defer w.mu.Unlock() + return append([]workerRequest(nil), w.requests...) +} + +var _ = Describe("The worker dialer", func() { + var ( + ctx context.Context + db *gorm.DB + clusterR *cluster.Registry + tunnels *tunnel.Registry + worker *scriptedWorker + ) + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + Expect(clusterR.Register(ctx, "replica-b", "test", 0, "")).To(Succeed()) + tunnels = tunnel.NewRegistry(clusterR, "replica-a") + worker = &scriptedWorker{} + }) + + // hold attaches a worker to this replica, with a bulk session when asked. + hold := func(bulk bool) (inference, bulkSess *tunnel.Session) { + GinkgoHelper() + front, back := sessionPair(tunnel.LaneInference) + worker.serve(back) + _, err := tunnels.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + if bulk { + bf, bb := sessionPair(tunnel.LaneBulk) + worker.serve(bb) + _, err = tunnels.Attach(ctx, "w1", tunnel.LaneBulk, bf) + Expect(err).ToNot(HaveOccurred()) + return front, bf + } + return front, nil + } + + roundTrip := func(conn net.Conn, msg string) { + GinkgoHelper() + _, err := conn.Write([]byte(msg)) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, len(msg)) + Expect(conn.SetReadDeadline(time.Now().Add(5 * time.Second))).To(Succeed()) + _, err = io.ReadFull(conn, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal(msg)) + } + + Describe("when this replica holds the tunnel", func() { + It("opens a stream down it, names the service, and hands back a conn that carries the protocol", func() { + hold(false) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "hello") + Expect(worker.seen()).To(Equal([]workerRequest{{tunnel.StreamTagGRPC, "10.0.0.5:50051"}})) + }) + + It("leaves no deadline on the conn, so that a quiet inference is not cut", func() { + hold(false) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + shortCtx, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer cancel() + conn, err := dialer.Dial(shortCtx, "w1", tunnel.StreamTagHTTP, "") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + time.Sleep(500 * time.Millisecond) + roundTrip(conn, "still here") + }) + + It("reports a refusal of the worker as the refusal of the worker, with no ErrNoRoute and no absence", func() { + hold(false) + worker.refuse = tunnel.ErrStreamTargetUnavailable + dialer := tunnel.NewWorkerDialer(tunnels, nil) + + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(err).To(HaveOccurred()) + Expect(tunnel.IsWorkerAnswer(err)).To(BeTrue()) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "a worker that answers is present") + Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse()) + }) + + It("reports that the worker learned nothing as ErrNoRoute and not as evidence", func() { + hold(false) + worker.refuse = tunnel.ErrStreamNotServed + dialer := tunnel.NewWorkerDialer(tunnels, nil) + + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("blames the budget of the caller and not the worker when the handshake ends on its deadline", func() { + hold(false) + worker.hold = 2 * time.Second + dialer := tunnel.NewWorkerDialer(tunnels, nil) + short, cancel := context.WithTimeout(ctx, 200*time.Millisecond) + defer cancel() + + _, err := dialer.Dial(short, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "an expired budget is not a missing route") + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("reports a session that ended under a held entry as a route that does not exist now", func() { + front, _ := hold(false) + Expect(front.Close()).To(Succeed()) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse(), "the tunnel is held here; looking elsewhere would come back to this replica") + }) + }) + + Describe("the bulk lane", func() { + It("carries a stream on the bulk session when the dial asks for it", func() { + hold(true) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagHTTP, "", tunnel.WithBulkLane()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "payload") + }) + + It("fails, and does not use the inference lane, when the worker has no bulk session", func() { + hold(false) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagHTTP, "", tunnel.WithBulkLane()) + Expect(errors.Is(err, tunnel.ErrNoBulkSession)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "a lane that is down must not make a scheduler demote a worker whose model calls work") + Expect(worker.seen()).To(BeEmpty(), "nothing may have crossed the inference lane") + }) + }) + + Describe("when another replica holds the tunnel", func() { + // The other replica is real: a registry of its own, a relay, and a worker + // behind it. The peer link between the two is a pair of pipes. + var ( + other *tunnel.Registry + peers *fakePeers + ) + + BeforeEach(func() { + other = tunnel.NewRegistry(clusterR, "replica-b") + front, back := sessionPair(tunnel.LaneInference) + worker.serve(back) + _, err := other.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + bf, bb := sessionPair(tunnel.LaneBulk) + worker.serve(bb) + _, err = other.Attach(ctx, "w1", tunnel.LaneBulk, bf) + Expect(err).ToNot(HaveOccurred()) + peers = &fakePeers{relay: tunnel.NewRelay(other), from: "replica-a"} + }) + + It("relays through the owner and carries bytes to the worker", func() { + dialer := tunnel.NewWorkerDialer(tunnels, peers) + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "through a peer") + Expect(peers.opened()).To(Equal([]string{"replica-b"})) + }) + + It("asks the owner for the bulk lane and for no fallback", func() { + dialer := tunnel.NewWorkerDialer(tunnels, peers) + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagHTTP, "", tunnel.WithBulkLane()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "file bytes") + Expect(peers.requested()).To(ContainElement(ContainSubstring("bulk-only"))) + }) + + It("gets the refusal of the owner when the owner has no bulk session, and not a transfer on the other lane", func() { + Expect(other.Disconnect("w1")).To(BeTrue()) + front, back := sessionPair(tunnel.LaneInference) + worker.serve(back) + _, err := other.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + + dialer := tunnel.NewWorkerDialer(tunnels, peers) + _, err = dialer.Dial(ctx, "w1", tunnel.StreamTagHTTP, "", tunnel.WithBulkLane()) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, tunnel.ErrRelayUnavailable)).To(BeTrue(), "%v", err) + Expect(worker.seen()).To(BeEmpty()) + }) + + It("states the remaining budget of the caller in the relay request", func() { + dialer := tunnel.NewWorkerDialer(tunnels, peers) + budgeted, cancel := context.WithTimeout(ctx, 20*time.Second) + defer cancel() + conn, err := dialer.Dial(budgeted, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + frames := peers.requested() + Expect(frames).To(HaveLen(1)) + Expect(frames[0]).To(MatchRegexp(`^w1 inference (1\d{4}|20000)$`), "a budget of about 20 seconds, in milliseconds") + }) + + It("states no budget for a caller that set no deadline", func() { + dialer := tunnel.NewWorkerDialer(tunnels, peers) + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Expect(peers.requested()).To(Equal([]string{"w1 inference 0"})) + }) + + It("passes a stale ownership refusal back as a routing fact", func() { + // The table names replica-b, but the registry behind its relay holds no + // session for the worker: the tunnel went away, and the row is a moment + // old. + peers.relay = tunnel.NewRelay(tunnel.NewRegistry(clusterR, "replica-b")) + dialer := tunnel.NewWorkerDialer(tunnels, peers) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeTrue(), "the cause is kept for a caller that can act on it") + Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse(), "absence must not reach the caller") + }) + + It("reports an unreachable peer as unreachable, with no ErrNoRoute and no absence", func() { + dialer := tunnel.NewWorkerDialer(tunnels, &downPeers{err: tunnelUnreachable()}) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "a link between two frontends says nothing about the worker") + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse()) + }) + + It("keeps an owner that was swept in the middle of the dial out of the chain", func() { + dialer := tunnel.NewWorkerDialer(tunnels, &downPeers{err: cluster.ErrInstanceNotFound}) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse()) + }) + + It("reports having no way to relay as its own condition", func() { + dialer := tunnel.NewWorkerDialer(tunnels, nil) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRelayPath)).To(BeTrue()) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeFalse(), "no peer was dialled") + }) + + It("does not relay to itself when the table names this replica", func() { + own := tunnel.NewRegistry(clusterR, "replica-a") + front, _ := sessionPair(tunnel.LaneInference) + _, err := own.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + // A registry that does not hold the session but whose replica owns the row. + dialer := tunnel.NewWorkerDialer(tunnel.NewRegistry(clusterR, "replica-a"), peers) + _, err = dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(peers.opened()).To(BeEmpty(), "relaying would send the request into this same process") + }) + }) + + Describe("when no live replica holds the tunnel", func() { + It("answers ErrNoRoute and never absence, for a worker with no connection row", func() { + dialer := tunnel.NewWorkerDialer(tunnels, &fakePeers{}) + _, err := dialer.Dial(ctx, "nobody", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse()) + }) + }) + + Describe("the dial functions it hands to the transports", func() { + It("binds one node and one tag, and passes the address through as the target", func() { + hold(false) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + conn, err := dialer.DialerFor("w1", tunnel.StreamTagHTTP)(ctx, "tcp", "w1.worker.invalid:80") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "x") + Expect(worker.seen()).To(Equal([]workerRequest{{tunnel.StreamTagHTTP, "w1.worker.invalid:80"}})) + }) + + It("gives gRPC a dialer fixed on the grpc tag", func() { + hold(false) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + conn, err := dialer.GRPCDialerFor("w1")(ctx, "10.0.0.5:50051") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "x") + Expect(worker.seen()).To(Equal([]workerRequest{{tunnel.StreamTagGRPC, "10.0.0.5:50051"}})) + }) + }) +}) + +// fakePeers opens a stream to the relay of another replica over a pipe, and +// records what the dialling side wrote first. +type fakePeers struct { + relay *tunnel.Relay + from string + + mu sync.Mutex + peers []string + frames []string +} + +func (f *fakePeers) Open(_ context.Context, peerID string) (net.Conn, error) { + f.mu.Lock() + f.peers = append(f.peers, peerID) + f.mu.Unlock() + if f.relay == nil { + return nil, errors.New("no relay behind this fake") + } + near, far := net.Pipe() + // The relay sees the far end. The near end is wrapped to record the frame + // that the dialler writes first. + go f.relay.Stream(f.from, far) + return &recordingConn{Conn: near, onFirstWrite: func(b []byte) { + f.mu.Lock() + defer f.mu.Unlock() + // A frame is a two byte length and the payload. + if len(b) > 2 { + f.frames = append(f.frames, string(b[2:])) + } + }}, nil +} + +func (f *fakePeers) opened() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.peers...) +} + +func (f *fakePeers) requested() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.frames...) +} + +type recordingConn struct { + net.Conn + once sync.Once + onFirstWrite func([]byte) +} + +func (c *recordingConn) Write(b []byte) (int, error) { + c.once.Do(func() { c.onFirstWrite(bytes.Clone(b)) }) + return c.Conn.Write(b) +} + +type downPeers struct{ err error } + +func (d *downPeers) Open(context.Context, string) (net.Conn, error) { return nil, d.err } + +// tunnelUnreachable returns an error that satisfies ErrPeerUnreachable the way +// the peer pool does: through a real pool dialling an address nobody listens on. +func tunnelUnreachable() error { + GinkgoHelper() + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg := cluster.NewRegistry(db) + Expect(reg.RegisterPeer(ctx, "ghost", "test", 0, "", "127.0.0.1:1", cluster.NewPeerCredential().Hash())).To(Succeed()) + pool := tunnel.NewPeerPool("replica-a", cluster.NewPeerCredential(), reg) + DeferCleanup(pool.Close) + _, err := pool.Open(ctx, "ghost") + Expect(err).To(MatchError(tunnel.ErrPeerUnreachable)) + return err +} diff --git a/core/services/tunnel/export_test.go b/core/services/tunnel/export_test.go index e70144f8e37d..c5d74c431280 100644 --- a/core/services/tunnel/export_test.go +++ b/core/services/tunnel/export_test.go @@ -16,3 +16,8 @@ func SetClaimTimeout(d time.Duration) (restore func()) { claimTimeout = d return func() { claimTimeout = old } } + +// NewRelayWithTimeouts builds a relay with timeouts that a spec can reach. +func NewRelayWithTimeouts(tunnels *Registry, header, open time.Duration) *Relay { + return newRelay(tunnels, header, open) +} diff --git a/core/services/tunnel/peerlink.go b/core/services/tunnel/peerlink.go new file mode 100644 index 000000000000..3eded8566ad7 --- /dev/null +++ b/core/services/tunnel/peerlink.go @@ -0,0 +1,320 @@ +package tunnel + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "net/url" + "sync" + "time" + + "github.com/gorilla/websocket" + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +// PeerPath is the route that a replica dials to open a peer link, and the route +// that the HTTP layer registers the handler on. The literal is spelled out here +// and not derived from the auth package, so that this package does not import +// it. A spec in the endpoints package, which sees both, keeps them equal. +const PeerPath = "/api/cluster/peer" + +// ErrPeerUnreachable reports that a peer which this deployment knows could not be +// reached: the dial failed, the peer refused the credentials, or its multiplexer +// would not carry a stream. +// +// It is not a form of ErrInstanceNotFound, and the two must stay apart. A caller +// that sees absence may conclude that a node is gone and reclaim what it ran. A +// caller that sees unreachability may only retry. It is not ErrNoRoute either: +// a link between two frontends is not a fact about any worker. See +// unreachableError for how this is enforced and not only documented. +var ErrPeerUnreachable = errors.New("tunnel: peer unreachable") + +// unreachableError reports a peer that could not be reached, with the cause in +// its message and out of its unwrap chain. The dial resolves the peer through the +// registry, so ErrInstanceNotFound is a cause this can be built over, and an +// unwrapped one would satisfy an absence check on a transport problem. +type unreachableError struct { + peerID string + cause error +} + +func (e *unreachableError) Error() string { + return fmt.Sprintf("tunnel: peer %q unreachable: %v", e.peerID, e.cause) +} + +// Unwrap reports only ErrPeerUnreachable. +func (e *unreachableError) Unwrap() error { return ErrPeerUnreachable } + +func unreachablePeer(peerID string, cause error) error { + return &unreachableError{peerID: peerID, cause: cause} +} + +// ErrPeerRejected reports that a peer answered and refused the credential of this +// replica. It has a different cure from the other ways a dial can end: a peer +// that does not answer is a transport problem that an operator waits out, and a +// peer that says no is a configuration problem. +// +// It is not absence. It says that a peer is up and talking. +var ErrPeerRejected = errors.New("tunnel: peer refused the credential of this replica") + +// rejectedError reports a peer that refused the dial. It unwraps to +// ErrPeerRejected and to ErrPeerUnreachable: the second keeps every caller that +// already handles an unreachable peer correct for a refusal, and the first lets a +// caller that wants to tell them apart do so. It unwraps to no absence sentinel. +type rejectedError struct { + peerID string + cause error +} + +func (e *rejectedError) Error() string { + return fmt.Sprintf("tunnel: peer %q refused the credential of this replica: %v", e.peerID, e.cause) +} + +func (e *rejectedError) Unwrap() []error { return []error{ErrPeerRejected, ErrPeerUnreachable} } + +func rejectedPeer(peerID string, cause error) error { + return &rejectedError{peerID: peerID, cause: cause} +} + +// ErrPoolClosed reports an Open on a pool that was shut down. It is a fact about +// this process and says nothing about whether the peer exists. +var ErrPoolClosed = errors.New("tunnel: peer pool is closed") + +// peerLinkHandshakeTimeout bounds the websocket upgrade of a peer link. It also +// bounds how long Close waits behind a dial in progress, which holds the lock of +// its peer. +const peerLinkHandshakeTimeout = 10 * time.Second + +// PeerPool dials peer replicas and keeps one multiplexed session for each. +// +// A peer link carries the traffic of every worker that the peer owns, so it is +// pooled: a dial for each relayed request would add a websocket handshake to +// every inference. +// +// The pool needs no knowledge of the errors of the multiplexer to keep its cache +// honest. A peer that shut down hands its session a go-away, and a session whose +// transport died hands out its shutdown error, and both arrive as a failed +// OpenStream that the same retry handles. A condition of one stream, such as a +// peer that reset one request, never reaches the pool, and dropping the session +// for it would tear down every other worker on the link. +type PeerPool struct { + selfID string + // cred is the own peer credential of this replica. It is what makes selfID a + // claim and not a label: the peer resolves selfID to a row and checks this + // against the hash that the row publishes. + cred PeerCredential + reg *cluster.Registry + + dialer *websocket.Dialer + + mu sync.Mutex + links map[string]*peerLink + closed bool +} + +// PeerCredential is the credential of a replica. It is cluster.PeerCredential; +// the alias keeps the callers of this package from importing both. +type PeerCredential = cluster.PeerCredential + +// peerLink is the cached session for one peer and the lock that serialises +// dialling it. The lock is per peer, so a dial to one peer that hangs does not +// hold up the opens to another. +type peerLink struct { + mu sync.Mutex + sess *Session +} + +// NewPeerPool returns a pool that dials peers as selfID. cred must be the value +// that the membership loop published the hash of, which makes a pool that +// presents one secret and a row that holds the hash of another unstateable. +// +// The pool authenticates with the credential of the replica and with nothing +// else. The registration token of the deployment opens no peer link, and an empty +// token therefore fails no peer dial. +func NewPeerPool(selfID string, cred PeerCredential, reg *cluster.Registry) *PeerPool { + return &PeerPool{ + selfID: selfID, + cred: cred, + reg: reg, + // No proxy: a peer link is inside one deployment, and honouring HTTP_PROXY + // would route it through whatever egress proxy the environment names. + dialer: NewDialer(peerLinkHandshakeTimeout), + links: map[string]*peerLink{}, + } +} + +// Open returns a stream to peerID, and dials and caches the session on first use. +// +// The errors are three conditions and callers act on them differently: +// ErrInstanceNotFound means the peer is not part of the deployment, +// ErrPeerUnreachable means it is and does not answer, and ErrPoolClosed means +// that this process is shutting down. Only the first is absence. +func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { + l, err := p.link(peerID) + if err != nil { + return nil, err + } + + l.mu.Lock() + defer l.mu.Unlock() + + if l.sess != nil { + st, err := l.sess.OpenStream(ctx) + if err == nil { + return st, nil + } + // A caller whose own budget ran out must not cost every other worker its + // link: the session is well, and this request is not. + if ctxErr := callerRanOut(ctx); ctxErr != nil { + return nil, ctxErr + } + // A session that died between calls is the common case. + xlog.Debug("tunnel peer link session unusable, dialling again", "peer", peerID, "error", err) + _ = l.sess.Close() + l.sess = nil + } + + sess, err := p.dial(ctx, peerID) + if err != nil { + // A dial that ran out of the time of the caller says nothing about the + // peer, which may be well. This also swallows a real ErrInstanceNotFound + // when the budget ran out in the same moment, which is the safe direction: + // a timeout must never manufacture absence. + if ctxErr := callerRanOut(ctx); ctxErr != nil { + return nil, ctxErr + } + return nil, err + } + + st, err := sess.OpenStream(ctx) + if err != nil { + // The peer answered and completed the handshake and will not carry a + // stream. That is a condition of the transport and never absence. + _ = sess.Close() + if ctxErr := callerRanOut(ctx); ctxErr != nil { + return nil, ctxErr + } + return nil, unreachablePeer(peerID, err) + } + l.sess = sess + return st, nil +} + +// link returns the entry for a peer and creates it on first use. Entries are not +// pruned: a peer that left the deployment but still listens keeps a websocket and +// two goroutines until Close or until the keepalive reclaims the session. +func (p *PeerPool) link(peerID string) (*peerLink, error) { + p.mu.Lock() + defer p.mu.Unlock() + if p.closed { + return nil, ErrPoolClosed + } + l, ok := p.links[peerID] + if !ok { + l = &peerLink{} + p.links[peerID] = l + } + return l, nil +} + +// dial resolves the advertised address of the peer and brings up the client side +// of a session over an authenticated websocket. +// +// A miss in the registry is returned as it is, so that ErrInstanceNotFound +// reaches the caller. Everything after it is wrapped as unreachable. The address +// is read only here: a peer that registers a new address while its session is +// alive keeps being reached over that session, which is right, because the peer +// still answers on the old one. +func (p *PeerPool) dial(ctx context.Context, peerID string) (*Session, error) { + inst, err := p.reg.Get(ctx, peerID) + if err != nil { + return nil, err + } + if inst.AdvertisedAddr == "" { + // A registered replica with no address is reachable by nobody. It is + // present, so this is not absence. + return nil, unreachablePeer(peerID, errors.New("peer has no advertised address")) + } + + endpoint := url.URL{ + // Plain ws: the link is authenticated by the credential of the replica and + // not by the transport. A deployment that wants the link encrypted puts the + // replicas behind TLS. + Scheme: "ws", + Host: inst.AdvertisedAddr, + Path: PeerPath, + RawQuery: url.Values{"id": []string{p.selfID}}.Encode(), + } + header := http.Header{} + // Sent even when empty. An empty value is refused by the peer, which is right + // for a process that never minted a credential, and omitting the header would + // make that indistinguishable from a release that predates it. + header.Set(cluster.PeerIdentityHeader, p.cred.Token()) + + ws, resp, err := p.dialer.DialContext(ctx, endpoint.String(), header) + status := 0 + if resp != nil { + status = resp.StatusCode + if resp.Body != nil { + _ = resp.Body.Close() + } + } + if err != nil { + // A peer that answered 401 is up and talking, and what it refused is the + // credential. The log line names the cure. + if status == http.StatusUnauthorized { + xlog.Warn("A peer refused the credential of this replica: the credential is not the one the peer has on record for this replica id. "+ + "This replica may have registered under an id that another process uses, or its row was replaced", + "peer", peerID, "self", p.selfID, "addr", inst.AdvertisedAddr) + return nil, rejectedPeer(peerID, err) + } + return nil, unreachablePeer(peerID, err) + } + + sess, err := peerClientSession(ws) + if err != nil { + _ = ws.Close() + return nil, unreachablePeer(peerID, err) + } + + // Close raced this dial. Handing the session back would leak it, since Close + // has already walked the map. + p.mu.Lock() + closed := p.closed + p.mu.Unlock() + if closed { + _ = sess.Close() + return nil, ErrPoolClosed + } + xlog.Debug("tunnel peer link dialled", "peer", peerID, "addr", inst.AdvertisedAddr) + return sess, nil +} + +// Close closes every cached session. It can be called twice, and an Open after it +// reports ErrPoolClosed. +func (p *PeerPool) Close() { + p.mu.Lock() + if p.closed { + p.mu.Unlock() + return + } + p.closed = true + links := p.links + p.links = nil + p.mu.Unlock() + + // Each session closes under the lock of its own peer, so that closing the pool + // cannot deadlock against an Open that is dialling and about to take p.mu. + for _, l := range links { + l.mu.Lock() + if l.sess != nil { + _ = l.sess.Close() + l.sess = nil + } + l.mu.Unlock() + } +} diff --git a/core/services/tunnel/peersessions.go b/core/services/tunnel/peersessions.go new file mode 100644 index 000000000000..f15d60bc0256 --- /dev/null +++ b/core/services/tunnel/peersessions.go @@ -0,0 +1,141 @@ +package tunnel + +import ( + "net" + "sync" + + "github.com/gorilla/websocket" + "github.com/mudler/xlog" +) + +// peerClientSession starts the dialling side of a peer link on ws. +// +// A peer link is a session with the windows of the bulk lane, because it carries +// the same kind of traffic, including the artifacts of models that are relayed to +// the owner of a tunnel. Both ends use the same configuration, since a window is +// announced by the side that receives. +// +// The dialling side accepts no stream. The far side opens none to it: it has a +// link of its own, dialled by itself, for what it sends to this replica. A +// stream that arrives here is reset, like a stream that a worker opens. +func peerClientSession(ws *websocket.Conn) (*Session, error) { + return clientSessionWith(ws, LaneBulk, 0) +} + +// PeerServerSession starts the accepting side of a peer link on ws. It accepts +// the streams of the dialling replica, up to the default limit of the +// multiplexer. A window is a bound on data received and not read, so a replica is +// sized against the windows times the limit for each link. +func PeerServerSession(ws *websocket.Conn) (*Session, error) { + return serverSessionWith(ws, LaneBulk, defaultIncomingStreams) +} + +// PeerSessions holds the peer links that this replica has accepted. It is the +// mirror of PeerPool: the pool owns the sessions that this replica dialled, and +// this owns those that its peers dialled into it. +// +// Something has to own an accepted session. The HTTP handler cannot, because it +// returns when the upgrade is done and the hijacked connection outlives it. And +// something has to accept the streams that arrive, because a stream is +// acknowledged only when the far side accepts it, so a session nobody accepts on +// does not fail an Open of the peer, it hangs it. +type PeerSessions struct { + // onStream handles one accepted stream and owns closing it. Nil closes the + // stream at once, which is what a replica with no relay does: it refuses + // promptly and does not leave a peer parked. + onStream func(peerID string, stream net.Conn) + + mu sync.Mutex + sessions map[string]*Session + closed bool +} + +// NewPeerSessions returns a holder whose accepted streams go to onStream. +func NewPeerSessions(onStream func(peerID string, stream net.Conn)) *PeerSessions { + return &PeerSessions{onStream: onStream, sessions: map[string]*Session{}} +} + +// Accept takes the ownership of a session that a peer dialled in, and returns at +// once: the loop that serves it runs on its own goroutine, because the return of +// the handler completes the hijack. +func (s *PeerSessions) Accept(peerID string, sess *Session) { + if sess == nil { + return + } + s.mu.Lock() + if s.closed { + s.mu.Unlock() + _ = sess.Close() + return + } + previous := s.sessions[peerID] + s.sessions[peerID] = sess + s.mu.Unlock() + + // A peer that dials again has lost its previous link, whether or not this side + // noticed. Keeping both would leave a session that nothing can be routed to. + if previous != nil { + xlog.Debug("tunnel peer dialled again, dropping its previous link", "peer", peerID) + _ = previous.Close() + } + go s.serve(peerID, sess) +} + +// Holds reports whether this replica holds an accepted link from peerID. A false +// answer is not an absent peer: it may be about to dial. +func (s *PeerSessions) Holds(peerID string) bool { + s.mu.Lock() + defer s.mu.Unlock() + sess, ok := s.sessions[peerID] + return ok && !sess.IsClosed() +} + +func (s *PeerSessions) serve(peerID string, sess *Session) { + defer func() { + s.forget(peerID, sess) + _ = sess.Close() + }() + for { + stream, err := sess.AcceptStream() + if err != nil { + // A link that ends is ordinary: a rolling update closes every session. + xlog.Debug("tunnel peer link ended", "peer", peerID, "error", err) + return + } + if s.onStream == nil { + xlog.Debug("tunnel peer stream refused: no relay installed", "peer", peerID) + _ = stream.Close() + continue + } + // One goroutine for each stream: the handler relays a whole request. + go s.onStream(peerID, stream) + } +} + +// forget drops the entry only if it still names this session. A peer that dialled +// again replaced it, and deleting blindly would evict the live link when the old +// one noticed that it died. +func (s *PeerSessions) forget(peerID string, sess *Session) { + s.mu.Lock() + defer s.mu.Unlock() + if s.sessions[peerID] == sess { + delete(s.sessions, peerID) + } +} + +// Close drops every held session. An Accept after it closes the session, so a +// dial that races the shutdown cannot leak a link. +func (s *PeerSessions) Close() { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return + } + s.closed = true + held := s.sessions + s.sessions = map[string]*Session{} + s.mu.Unlock() + for _, sess := range held { + _ = sess.Close() + } +} diff --git a/core/services/tunnel/relay.go b/core/services/tunnel/relay.go new file mode 100644 index 000000000000..dd44228a9e09 --- /dev/null +++ b/core/services/tunnel/relay.go @@ -0,0 +1,361 @@ +package tunnel + +import ( + "cmp" + "context" + "errors" + "fmt" + "io" + "net" + "strconv" + "strings" + "time" + + "github.com/mudler/xlog" +) + +// The framing that every stream of a peer link opens with. +// +// A worker holds one tunnel and it lands on one frontend replica, so every other +// replica reaches that worker through the one that holds it. A stream of a peer +// link carries the traffic of every worker that the far replica owns, so it means +// nothing until it says which worker it is for. That is this frame. +// +// A relayed stream carries two request frames one after the other. The first is +// this one, and the owning replica consumes it. The second is the frame of the +// worker tunnel, which crosses untouched and is answered by the worker. A dialler +// reads one reply from each, in that order. +// +// The two vocabularies are disjoint on purpose. "relay-ok" is not "ok", and no +// refusal code below is spelled like a code of the tunnel, so a reader applied to +// the wrong hop fails with "unrecognised reply" and does not hand back a +// plausible sentinel of the other hop. Reporting a refusal of the worker as one of +// the owning replica would send the caller to retry against the wrong end. +const ( + relayReplyAccepted = "relay-ok" + relayCodeNotOwner = "relay-not-owner" + relayCodeUnavailable = "relay-unavailable" + relayCodeBadRequest = "relay-bad-request" +) + +// The lanes a relay request can name. The third keeps the refusal to fall back: +// the dialler asked for the bulk lane and for no other. +const ( + relayLaneInference = "inference" + relayLaneBulk = "bulk" + relayLaneBulkOnly = "bulk-only" +) + +func relayLane(o dialOptions) string { + switch { + case o.lane == LaneBulk && o.noFallback: + return relayLaneBulkOnly + case o.lane == LaneBulk: + return relayLaneBulk + } + return relayLaneInference +} + +// The refusals that the relay hop can send, beyond ErrNotOwner, which it shares +// with the local path. +// +// They are kept apart. ErrNotOwner is a routing fact: the worker may be well on +// another replica, and the caller resolves the owner again. ErrRelayUnavailable +// is the owning replica failing: it holds the tunnel and its session will not +// carry a stream, so a retry is worth something and looking elsewhere is not. +// ErrRelayRequestInvalid is a bug of the caller, and no retry helps. +// +// None of them is built over an absence error. A refusal proves that a replica +// answered. +var ( + ErrRelayUnavailable = errors.New("tunnel: the owning replica could not open a stream to that worker") + ErrRelayRequestInvalid = errors.New("tunnel: the owning replica rejected the relay request as malformed") +) + +// WriteRelayRequest names the worker that a peer stream is for, the lane, and how +// much time the original caller still has. +// +// The budget is what makes the open of the relay honest. Everything past this +// frame is work for a caller that the relay cannot see, so without a budget the +// relay can only use a constant that no operator can set. Zero means that the +// budget is not stated, and it is written as 0. A caller with no time left is +// about to fail on its own context, and stating less than a millisecond is +// written as 1. +func WriteRelayRequest(w io.Writer, nodeID, lane string, budget time.Duration) error { + if nodeID == "" { + return errors.New("writing a relay request: empty node id") + } + if strings.Contains(nodeID, streamRequestSeparator) { + return fmt.Errorf("writing a relay request: node id %q contains a space", nodeID) + } + switch lane { + case relayLaneInference, relayLaneBulk, relayLaneBulkOnly: + default: + return fmt.Errorf("writing a relay request: unknown lane %q", lane) + } + var millis int64 + if budget > 0 { + // A count of milliseconds has one spelling. It is rounded up, so that a + // sub-millisecond budget stays positive and keeps meaning "almost none". + millis = int64((budget + time.Millisecond - 1) / time.Millisecond) + } + return writeFrame(w, nodeID+streamRequestSeparator+lane+streamRequestSeparator+strconv.FormatInt(millis, 10)) +} + +// ReadRelayRequest reads the opening frame of a peer stream. The budget is zero +// when the dialling replica stated none. +// +// A malformed frame is an ordinary error and not ErrRelayRequestInvalid. That +// sentinel is what a relay sends to describe a refusal, and returning it here +// would leave a caller unable to tell "the peer refused my request" from "I could +// not read the request of the peer". +func ReadRelayRequest(r io.Reader) (nodeID, lane string, budget time.Duration, err error) { + payload, err := readFrame(r) + if err != nil { + return "", "", 0, fmt.Errorf("reading a relay request: %w", err) + } + fields := strings.Split(payload, streamRequestSeparator) + if len(fields) != 3 || fields[0] == "" { + return "", "", 0, fmt.Errorf("reading a relay request: want \"node lane budget\", got %d fields", len(fields)) + } + nodeID, lane = fields[0], fields[1] + switch lane { + case relayLaneInference, relayLaneBulk, relayLaneBulkOnly: + default: + return "", "", 0, fmt.Errorf("reading a relay request for node %q: unknown lane %q", nodeID, lane) + } + millis, err := strconv.ParseInt(fields[2], 10, 64) + if err != nil || millis < 0 { + return "", "", 0, fmt.Errorf("reading a relay request for node %q: budget %q is not a number of milliseconds", nodeID, fields[2]) + } + // time.Duration is nanoseconds in an int64, so a large count overflows when it + // is multiplied. A caller that claims to wait longer than the backstop of the + // relay gets the backstop either way, so the count is clamped. + millis = min(millis, maxRelayBudgetMillis) + return nodeID, lane, time.Duration(millis) * time.Millisecond, nil +} + +// WriteRelayAccepted tells the peer that the stream now carries the conversation +// of the worker tunnel. Everything after this frame belongs to that hop. +func WriteRelayAccepted(w io.Writer) error { return writeFrame(w, relayReplyAccepted) } + +// WriteRelayRefusal reports why a peer stream is not relayed. The caller closes +// the stream after it. A reason that has no class is sent as bad-request with its +// text, and not dropped: a refusal that a peer cannot read looks like a replica +// that hung up. +func WriteRelayRefusal(w io.Writer, reason error) error { + code := relayCodeBadRequest + switch { + case errors.Is(reason, ErrNotOwner): + code = relayCodeNotOwner + case errors.Is(reason, ErrRelayUnavailable): + code = relayCodeUnavailable + } + text := "" + if reason != nil { + text = strings.Map(func(r rune) rune { + if r == '\n' || r == '\r' { + return ' ' + } + return r + }, reason.Error()) + } + return writeFrame(w, truncateRunes(replyPrefixRefused+code+streamRequestSeparator+text, maxFrame)) +} + +// ReadRelayReply reads the answer of the owning replica. A nil result means that +// the stream is now the stream of the worker tunnel. +// +// A failure to read the reply is returned as itself and never as one of the +// refusals. A refusal means that a replica answered, and a failed read means that +// the peer link broke. Reporting the second as the first would present a dead +// link as a decision. +func ReadRelayReply(r io.Reader) error { + payload, err := readFrame(r) + if err != nil { + return fmt.Errorf("reading a relay reply: %w", err) + } + if payload == relayReplyAccepted { + return nil + } + rest, ok := strings.CutPrefix(payload, replyPrefixRefused) + if !ok { + return fmt.Errorf("reading a relay reply: unrecognised reply %q", payload) + } + code, text, _ := strings.Cut(rest, streamRequestSeparator) + switch code { + case relayCodeNotOwner: + // ErrNotOwner and nothing else: a routing fact. It is not ErrNoConnection + // (the worker is connected nowhere) and not ErrPeerUnreachable (a replica + // will not answer). + return fmt.Errorf("%w: %s", ErrNotOwner, text) + case relayCodeUnavailable: + return fmt.Errorf("%w: %s", ErrRelayUnavailable, text) + case relayCodeBadRequest: + return fmt.Errorf("%w: %s", ErrRelayRequestInvalid, text) + } + // A code of a newer replica. It is not mapped onto the nearest known code, so + // that a caller does not retry for ever against a refusal that means something + // else. + return fmt.Errorf("relay stream refused with unrecognised code %q: %s", code, text) +} + +// maxRelayBudgetMillis is the largest budget a peer may declare. A day is far +// past relayOpenTimeout, which is all a budget is compared with. +const maxRelayBudgetMillis = int64(24 * 60 * 60 * 1000) + +const ( + // relayHeaderTimeout bounds how long a peer stream may go without naming the + // worker it is for. Without it, a dialler that is killed between OpenStream + // and its first write holds a goroutine and a slot of the link until the link + // dies. + relayHeaderTimeout = 15 * time.Second + + // relayOpenTimeout is the ceiling on opening the stream to the worker. yamux + // blocks an open once its accept backlog is full, and honours the context, + // which is the reason there is one. It is the backstop for a caller that + // stated no budget, and the stated budget only ever shortens it: a caller + // willing to wait an hour must not be able to park a goroutine and a slot on a + // worker that stopped accepting. + relayOpenTimeout = 15 * time.Second +) + +// Relay splices a stream that a peer opened onto a worker tunnel that this +// replica holds. +// +// It is what makes more than one replica work. With N replicas behind a load +// balancer, (N-1)/N of the requests arrive at a replica that does not hold the +// tunnel, and they reach the worker through here. +// +// One hop, always. A stream that names a worker this replica does not hold is +// refused and never relayed onward. A second hop would turn a stale ownership row +// into a loop between two replicas, each sure that the other holds the worker. +type Relay struct { + tunnels *Registry + + // The timeouts are fields so that a spec can reach a deadline without waiting + // out a production value. + headerTimeout time.Duration + openTimeout time.Duration +} + +// NewRelay returns the relay for the tunnels that this replica holds. Its Stream +// method is the handler of a PeerSessions. +func NewRelay(tunnels *Registry) *Relay { return newRelay(tunnels, 0, 0) } + +func newRelay(tunnels *Registry, headerTimeout, openTimeout time.Duration) *Relay { + return &Relay{ + tunnels: tunnels, + headerTimeout: cmp.Or(headerTimeout, relayHeaderTimeout), + openTimeout: cmp.Or(openTimeout, relayOpenTimeout), + } +} + +// Stream relays one peer stream. It owns closing the stream on every path. +func (r *Relay) Stream(peerID string, stream net.Conn) { + // It runs on a bare goroutine, and a panic that is not recovered ends the + // process, with the traffic of every other replica that goes through this one. + // The recovery covers the frame read, the lookup and the open. It cannot cover + // the copy goroutines of Splice, and it does not panic again, because no + // middleware sits above a goroutine that the HTTP layer already left. + defer func() { + if p := recover(); p != nil { + xlog.Error("Panic while relaying a peer stream", "peer", peerID, "panic", p) + _ = stream.Close() + } + }() + + local, ok := r.accept(peerID, stream) + if !ok { + return + } + // Splice owns closing both ends. A relayed request that a client abandons + // ends with an error, so it is logged at debug level only. + if err := Splice(stream, local); err != nil { + xlog.Debug("relayed peer stream ended with an error", "peer", peerID, "error", err) + } +} + +// accept reads which worker the stream is for and opens the stream to the worker. +// The second result is false when the stream was refused, and then the refusal +// was sent and the stream closed. +func (r *Relay) accept(peerID string, stream net.Conn) (net.Conn, bool) { + if err := stream.SetReadDeadline(time.Now().Add(r.headerTimeout)); err != nil { + r.refuse(peerID, stream, fmt.Errorf("%w: arming the request deadline: %v", ErrRelayUnavailable, err)) + return nil, false + } + + nodeID, lane, budget, err := ReadRelayRequest(stream) + if err != nil { + // This includes the deadline above running out. Both mean that the stream + // never said which worker it wanted, which is a bug of the dialling + // replica, and a retry against this one does not cure it. + r.refuse(peerID, stream, fmt.Errorf("%w: %v", ErrRelayRequestInvalid, err)) + return nil, false + } + + // Cleared before the open, and nothing arms another deadline after it. What + // follows is a relayed request whose length is the business of the caller, and + // a deadline left armed would end a long inference after a quiet moment. The + // keepalive of the peer link kills a session whose far side stopped answering. + if err := stream.SetReadDeadline(time.Time{}); err != nil { + r.refuse(peerID, stream, fmt.Errorf("%w: clearing the request deadline: %v", ErrRelayUnavailable, err)) + return nil, false + } + + // The smaller of the ceiling and what the caller says it still has. + open := r.openTimeout + if budget > 0 && budget < open { + open = budget + } + ctx, cancel := context.WithTimeout(context.Background(), open) + defer cancel() + + var ( + openLane = LaneInference + openOpts []OpenOption + ) + switch lane { + case relayLaneBulk: + openLane = LaneBulk + case relayLaneBulkOnly: + openLane = LaneBulk + openOpts = append(openOpts, WithoutFallback()) + } + local, err := r.tunnels.Open(ctx, nodeID, openLane, openOpts...) + if err != nil { + if errors.Is(err, ErrNotOwner) { + // As itself: the worker is probably well on another replica, and this + // is the answer that tells the caller to look there. + r.refuse(peerID, stream, err) + return nil, false + } + // Everything else is this replica failing. It must not become ErrNotOwner, + // because the tunnel is held here and the caller would come back to this + // replica. It must not become absence either: the worker is attached. + r.refuse(peerID, stream, fmt.Errorf("%w: %v", ErrRelayUnavailable, err)) + return nil, false + } + + if err := WriteRelayAccepted(stream); err != nil { + // The peer never learns that the stream was accepted, so it cannot use it. + // The stream to the worker is closed, or one would leak for every failed + // reply. + xlog.Debug("could not accept a peer stream for relaying", "peer", peerID, "node", nodeID, "error", err) + _ = local.Close() + _ = stream.Close() + return nil, false + } + return local, true +} + +// refuse reports why a stream is not relayed and then ends it. The close is not +// optional: a replica that says why and leaves the stream open has parked the +// peer on a request that will never be served, and that looks like a slow replica +// and not like a refusal. The reply is best effort and the close is not. +func (r *Relay) refuse(peerID string, stream net.Conn, reason error) { + if err := WriteRelayRefusal(stream, reason); err != nil { + xlog.Debug("could not tell a peer why its stream was refused", "peer", peerID, "reason", reason, "error", err) + } + _ = stream.Close() +} diff --git a/core/services/tunnel/relay_test.go b/core/services/tunnel/relay_test.go new file mode 100644 index 000000000000..3c23efa5d6b8 --- /dev/null +++ b/core/services/tunnel/relay_test.go @@ -0,0 +1,265 @@ +package tunnel_test + +import ( + "bytes" + "context" + "errors" + "net" + "strings" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("The relay request frame", func() { + It("round-trips the node, the lane and a stated budget", func() { + var buf bytes.Buffer + Expect(tunnel.WriteRelayRequest(&buf, "w1", "bulk", 1500*time.Millisecond)).To(Succeed()) + node, lane, budget, err := tunnel.ReadRelayRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(node).To(Equal("w1")) + Expect(lane).To(Equal("bulk")) + Expect(budget).To(Equal(1500 * time.Millisecond)) + }) + + It("writes zero for a budget that is not stated, and reads it back as not stated", func() { + var buf bytes.Buffer + Expect(tunnel.WriteRelayRequest(&buf, "w1", "inference", 0)).To(Succeed()) + _, _, budget, err := tunnel.ReadRelayRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(budget).To(BeZero()) + }) + + It("rounds a sub-millisecond budget up, so that it keeps meaning almost none", func() { + var buf bytes.Buffer + Expect(tunnel.WriteRelayRequest(&buf, "w1", "inference", time.Microsecond)).To(Succeed()) + _, _, budget, err := tunnel.ReadRelayRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(budget).To(Equal(time.Millisecond)) + }) + + It("refuses a node id that would split across the separator, an empty one, and an unknown lane", func() { + var buf bytes.Buffer + Expect(tunnel.WriteRelayRequest(&buf, "w 1", "inference", 0)).ToNot(Succeed()) + Expect(tunnel.WriteRelayRequest(&buf, "", "inference", 0)).ToNot(Succeed()) + Expect(tunnel.WriteRelayRequest(&buf, "w1", "express", 0)).ToNot(Succeed()) + Expect(buf.Len()).To(BeZero(), "nothing may reach the wire") + }) + + DescribeTable("rejects a frame it cannot read, as an ordinary error", + func(payload string) { + var buf bytes.Buffer + frame := append([]byte{byte(len(payload) >> 8), byte(len(payload))}, payload...) + buf.Write(frame) + _, _, _, err := tunnel.ReadRelayRequest(&buf) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, tunnel.ErrRelayRequestInvalid)).To(BeFalse(), + "that sentinel is what a relay sends; a reader that returned it would confuse the two ends") + }, + Entry("a node alone", "w1"), + Entry("a missing budget", "w1 bulk"), + Entry("a budget that is not a number", "w1 bulk soon"), + Entry("a negative budget", "w1 bulk -5"), + Entry("an unknown lane", "w1 express 0"), + Entry("no node", " bulk 0"), + Entry("too many fields", "w1 bulk 0 extra"), + ) + + It("clamps a budget above a day, so that the multiplication cannot overflow", func() { + var buf bytes.Buffer + payload := "w1 inference 9223372036854775807" + buf.Write(append([]byte{byte(len(payload) >> 8), byte(len(payload))}, payload...)) + _, _, budget, err := tunnel.ReadRelayRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(budget).To(Equal(24 * time.Hour)) + }) +}) + +var _ = Describe("The relay reply", func() { + roundTrip := func(reason error) error { + var buf bytes.Buffer + Expect(tunnel.WriteRelayRefusal(&buf, reason)).To(Succeed()) + return tunnel.ReadRelayReply(&buf) + } + + It("reads the acceptance as nil", func() { + var buf bytes.Buffer + Expect(tunnel.WriteRelayAccepted(&buf)).To(Succeed()) + Expect(tunnel.ReadRelayReply(&buf)).To(Succeed()) + }) + + It("keeps the three refusals apart", func() { + Expect(roundTrip(tunnel.ErrNotOwner)).To(MatchError(tunnel.ErrNotOwner)) + Expect(roundTrip(tunnel.ErrRelayUnavailable)).To(MatchError(tunnel.ErrRelayUnavailable)) + Expect(roundTrip(errors.New("anything else"))).To(MatchError(tunnel.ErrRelayRequestInvalid)) + }) + + It("never reads a refusal of the owner as one of the worker, nor the reverse", func() { + Expect(tunnel.IsWorkerAnswer(roundTrip(tunnel.ErrNotOwner))).To(BeFalse()) + Expect(tunnel.IsWorkerAnswer(roundTrip(tunnel.ErrRelayUnavailable))).To(BeFalse()) + + var buf bytes.Buffer + Expect(tunnel.WriteStreamRefusal(&buf, tunnel.ErrStreamTargetUnavailable)).To(Succeed()) + err := tunnel.ReadRelayReply(&buf) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, tunnel.ErrStreamTargetUnavailable)).To(BeFalse(), "a reader of the wrong hop fails with an unrecognised reply") + }) + + It("reports a code of a newer replica as itself and not as the nearest known code", func() { + payload := "err relay-teleport somewhere else" + var buf bytes.Buffer + buf.Write(append([]byte{byte(len(payload) >> 8), byte(len(payload))}, payload...)) + err := tunnel.ReadRelayReply(&buf) + Expect(err).To(MatchError(ContainSubstring("unrecognised code"))) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse()) + Expect(errors.Is(err, tunnel.ErrRelayUnavailable)).To(BeFalse()) + }) + + It("reports a link that broke as a failed read and not as a refusal", func() { + err := tunnel.ReadRelayReply(strings.NewReader("")) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse()) + Expect(errors.Is(err, tunnel.ErrRelayUnavailable)).To(BeFalse()) + }) +}) + +var _ = Describe("The relay", func() { + var ( + ctx context.Context + clusterR *cluster.Registry + tunnels *tunnel.Registry + worker *scriptedWorker + ) + + BeforeEach(func() { + ctx = context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-b", "test", 0, "")).To(Succeed()) + tunnels = tunnel.NewRegistry(clusterR, "replica-b") + worker = &scriptedWorker{} + front, back := sessionPair(tunnel.LaneInference) + worker.serve(back) + _, err := tunnels.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + }) + + // ask sends one request frame to a relay and returns the reply. + ask := func(frame func(net.Conn)) (net.Conn, error) { + near, far := net.Pipe() + DeferCleanup(func() { _ = near.Close() }) + go tunnel.NewRelay(tunnels).Stream("replica-a", far) + frame(near) + _ = near.SetReadDeadline(time.Now().Add(5 * time.Second)) + return near, tunnel.ReadRelayReply(near) + } + + It("relays a stream to the worker it names and splices the two", func() { + near, err := ask(func(c net.Conn) { + Expect(tunnel.WriteRelayRequest(c, "w1", "inference", 0)).To(Succeed()) + }) + Expect(err).ToNot(HaveOccurred()) + Expect(tunnel.WriteStreamRequest(near, tunnel.StreamTagGRPC, "x:1")).To(Succeed()) + Expect(tunnel.ReadStreamReply(near)).To(Succeed()) + }) + + It("refuses a worker it does not hold with ErrNotOwner, and does not relay onward", func() { + _, err := ask(func(c net.Conn) { + Expect(tunnel.WriteRelayRequest(c, "w-elsewhere", "inference", 0)).To(Succeed()) + }) + Expect(err).To(MatchError(tunnel.ErrNotOwner)) + }) + + It("refuses a request it cannot read as a bad request", func() { + _, err := ask(func(c net.Conn) { _, _ = c.Write([]byte{0, 3, 'a', 'b', 'c'}) }) + Expect(err).To(MatchError(tunnel.ErrRelayRequestInvalid)) + }) + + It("refuses with ErrRelayUnavailable when the tunnel is held and will not carry a stream", func() { + Expect(tunnels.Disconnect("w1")).To(BeTrue()) + front, _ := sessionPair(tunnel.LaneInference) + _, err := tunnels.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + Expect(front.Close()).To(Succeed()) + + _, err = ask(func(c net.Conn) { + Expect(tunnel.WriteRelayRequest(c, "w1", "inference", 0)).To(Succeed()) + }) + Expect(err).To(MatchError(tunnel.ErrRelayUnavailable)) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeFalse(), "the tunnel is held here; the caller would come back to this replica") + }) + + It("ends a stream that never says which worker it is for", func() { + near, far := net.Pipe() + DeferCleanup(func() { _ = near.Close() }) + done := make(chan struct{}) + relay := tunnel.NewRelayWithTimeouts(tunnels, 150*time.Millisecond, time.Second) + go func() { relay.Stream("replica-a", far); close(done) }() + _ = near.SetReadDeadline(time.Now().Add(5 * time.Second)) + err := tunnel.ReadRelayReply(near) + Expect(err).To(MatchError(tunnel.ErrRelayRequestInvalid)) + Eventually(done).Should(BeClosed()) + }) + + It("takes the smaller of its ceiling and the stated budget for the open", func() { + // A worker that holds a session which never accepts a stream: the open + // blocks on the backlog. The stated budget ends it long before the ceiling. + Expect(tunnels.Disconnect("w1")).To(BeTrue()) + front, _ := sessionPair(tunnel.LaneInference) + _, err := tunnels.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + // Fill the accept backlog of the far side, which accepts nothing. + var ( + heldMu sync.Mutex + held []net.Conn + ) + DeferCleanup(func() { + heldMu.Lock() + defer heldMu.Unlock() + for _, c := range held { + _ = c.Close() + } + }) + for range 300 { + go func() { + c, err := tunnels.Open(ctx, "w1", tunnel.LaneInference) + if err == nil { + heldMu.Lock() + held = append(held, c) + heldMu.Unlock() + } + }() + } + time.Sleep(300 * time.Millisecond) + + begun := time.Now() + _, rerr := ask(func(c net.Conn) { + Expect(tunnel.WriteRelayRequest(c, "w1", "inference", 300*time.Millisecond)).To(Succeed()) + }) + Expect(rerr).To(MatchError(tunnel.ErrRelayUnavailable)) + Expect(time.Since(begun)).To(BeNumerically("<", 5*time.Second)) + }) +}) + +var _ = Describe("Route failures", func() { + It("never lets an absence claim out, and keeps ErrNoRoute on a route that does not exist", func() { + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg := cluster.NewRegistry(db) + Expect(reg.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + dialer := tunnel.NewWorkerDialer(tunnel.NewRegistry(reg, "replica-a"), &downPeers{err: cluster.ErrInstanceNotFound}) + _, err := dialer.Dial(ctx, "nobody", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue()) + Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse()) + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse()) + Expect(err.Error()).To(ContainSubstring("nobody")) + }) +}) diff --git a/core/services/tunnel/route.go b/core/services/tunnel/route.go new file mode 100644 index 000000000000..6164cf06d904 --- /dev/null +++ b/core/services/tunnel/route.go @@ -0,0 +1,87 @@ +package tunnel + +import ( + "context" + "errors" + "fmt" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +// ErrNoRoute reports that this replica could not get a request to a worker right +// now. It is a routing fact and nothing more. +// +// It says nothing about whether the worker exists or runs. A worker that is +// registered, heartbeating and serving can be unroutable from here for ordinary +// reasons: it has not dialled its tunnel yet after an upgrade of the frontends, +// the replica that holds its tunnel is restarting, the ownership row is a moment +// old, or this replica has no way to reach the owner. +// +// Every failure to resolve or open a route carries it, so that a consumer that +// must not act on absence has one check to make. The specific cause stays in the +// chain for a caller that can act on it, with the exception that +// routeFailure describes. +// +// It is not carried by: +// - a refusal that the worker wrote, because a worker that answers is present; +// - a peer replica that did not answer (ErrPeerUnreachable), because that is a +// fact about the link between two frontends and a scheduler that demoted a +// worker for it would take capacity away for no reason of the worker; +// - the budget of the caller running out, because an expiry says nothing about +// any worker. +var ErrNoRoute = errors.New("tunnel: no route from this replica to that worker") + +// ErrNoRelayPath reports that the tunnel of a worker is held by another replica +// and this one has no way to reach that replica. It is its own condition: not +// ErrNotOwner, which tells a caller to resolve the owner again, and not +// ErrPeerUnreachable, because no peer was dialled. +var ErrNoRelayPath = errors.New("tunnel: this replica cannot relay to the owner of that worker") + +// noRouteError reports a worker that this replica cannot route to, keeping the +// cause in its message and out of its unwrap chain. +// +// The causes are absence claims: no live replica holds the tunnel of the worker, +// and no such replica is registered. Both are true statements about the cluster. +// Neither is a statement about the worker, and a consumer that matched on them +// would conclude that a worker that is heartbeating and serving has gone away and +// reclaim its models. The guarantee belongs to the type: no call site can leak an +// absence sentinel through it. +type noRouteError struct { + nodeID string + cause error +} + +func (e *noRouteError) Error() string { + return fmt.Sprintf("tunnel: no route from this replica to node %q: %v", e.nodeID, e.cause) +} + +// Unwrap reports only ErrNoRoute. +func (e *noRouteError) Unwrap() error { return ErrNoRoute } + +// isAbsenceClaim reports whether an error asserts that something does not exist. +func isAbsenceClaim(err error) bool { + return errors.Is(err, cluster.ErrNoConnection) || errors.Is(err, cluster.ErrInstanceNotFound) +} + +// routeFailure turns the failure of a dial into the error that the caller gets. +// It is the one place that decides, so that two predicates cannot drift apart. +// +// - An absence claim does not reach the caller. It becomes ErrNoRoute with the +// cause in the message only. +// - A failure of the budget of the caller, of the link to a peer replica, of the +// pool, or of the bulk lane alone is returned as itself, wrapped, with no +// ErrNoRoute. None of them is a statement about whether the worker has a +// route, and ErrNoRoute makes a scheduler demote the worker. +// - Everything else is a route that does not exist now: ErrNoRoute on top, and +// the cause below it. +func routeFailure(nodeID string, cause error) error { + switch { + case isAbsenceClaim(cause): + return &noRouteError{nodeID: nodeID, cause: cause} + case errors.Is(cause, context.DeadlineExceeded), errors.Is(cause, context.Canceled), + errors.Is(cause, ErrPeerUnreachable), errors.Is(cause, ErrPoolClosed), + errors.Is(cause, ErrNoBulkSession): + return fmt.Errorf("reaching node %q: %w", nodeID, cause) + } + return fmt.Errorf("reaching node %q: %w: %w", nodeID, ErrNoRoute, cause) +} diff --git a/core/services/tunnel/session.go b/core/services/tunnel/session.go index 9cc9f8e2ff0f..ab32fe5e5cdc 100644 --- a/core/services/tunnel/session.go +++ b/core/services/tunnel/session.go @@ -90,8 +90,18 @@ type Session struct { // misbehaves. The worker side keeps the default, because it accepts the // streams of the frontend. func ServerSession(ws *websocket.Conn, lane Lane) (*Session, error) { + return serverSessionWith(ws, lane, 0) +} + +// defaultIncomingStreams is the limit of yamux on the streams that the other side +// may have open. +const defaultIncomingStreams = 1000 + +// serverSessionWith starts the server side of a session that accepts at most +// incoming streams opened by the other side. Zero resets every one of them. +func serverSessionWith(ws *websocket.Conn, lane Lane, incoming uint32) (*Session, error) { cfg := sessionConfig(lane) - cfg.MaxIncomingStreams = 0 + cfg.MaxIncomingStreams = incoming s, err := yamux.Server(websocketConn(ws), cfg, nil) if err != nil { return nil, fmt.Errorf("starting the server side of a tunnel session: %w", err) @@ -102,7 +112,13 @@ func ServerSession(ws *websocket.Conn, lane Lane) (*Session, error) { // ClientSession starts the worker side of a session on ws. The worker accepts // the streams that the frontend opens. func ClientSession(ws *websocket.Conn, lane Lane) (*Session, error) { - s, err := yamux.Client(websocketConn(ws), sessionConfig(lane), nil) + return clientSessionWith(ws, lane, defaultIncomingStreams) +} + +func clientSessionWith(ws *websocket.Conn, lane Lane, incoming uint32) (*Session, error) { + cfg := sessionConfig(lane) + cfg.MaxIncomingStreams = incoming + s, err := yamux.Client(websocketConn(ws), cfg, nil) if err != nil { return nil, fmt.Errorf("starting the client side of a tunnel session: %w", err) } From e33662df3946003ef3f340933bc44da4d14b17ee Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 00:53:05 +0000 Subject: [PATCH 059/162] feat(http): serve the link that replicas dial to reach a tunnel GET /api/cluster/peer upgrades an authenticated replica to a link and hands the session to the holder, which relays its streams onto the tunnels this replica holds. One credential decides: the secret of the dialling replica against the hash in its instances row. The registration token of the deployment opens no link, so a worker that holds it cannot pose as a replica. The credential is checked before the upgrade, and the answers are 401 for a dial with none, an unknown replica, an empty hash or a wrong secret, and 503 where there is no table of replicas. The route is registered in every deployment so that the route coverage spec sees it. LOCALAI_PEER_ADDRESS sets the address that peers dial. Without it the address is found from the route to the database. A replica with no usable address still starts and publishes none. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 76 ++++- core/cli/run.go | 4 + core/config/distributed_config.go | 16 +- core/http/app.go | 11 +- core/http/auth/helpers_test.go | 2 +- core/http/auth/middleware_test.go | 17 ++ core/http/auth/public_routes.go | 7 +- core/http/endpoints/cluster/peer.go | 122 +++++++++ core/http/endpoints/cluster/peer_test.go | 305 +++++++++++++++++++++ core/http/routes/cluster.go | 9 +- core/services/nodes/control_tunnel_test.go | 8 +- 11 files changed, 566 insertions(+), 11 deletions(-) create mode 100644 core/http/endpoints/cluster/peer.go create mode 100644 core/http/endpoints/cluster/peer_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index ee3cb42ff149..22309eaf724a 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "io" + "net" + "strconv" "strings" "sync" "sync/atomic" @@ -77,6 +79,19 @@ type DistributedServices struct { // every replica and refuses every worker that has no tunnel credential. Tunnels *tunnel.Registry + // Instances reads the instances table. The peer route checks the + // credential of a dialling replica against it. + Instances *cluster.Registry + // PeerSessions holds the links that other replicas dialled into this one, + // and relays their streams onto the tunnels that this replica holds. + PeerSessions *tunnel.PeerSessions + // PeerPool holds the links that this replica dialled to its peers. + PeerPool *tunnel.PeerPool + // WorkerDialer opens a stream to a worker through the replica that holds its + // tunnel: directly, or through a peer link. The tunnel carrier builds the + // dial seams on it; nothing calls it while the cluster runs on NATS. + WorkerDialer *tunnel.WorkerDialer + // active names the carrier set the holders forward to. active *atomic.Pointer[carrier.Set] @@ -95,6 +110,12 @@ func (ds *DistributedServices) Shutdown() { if ds.Tunnels != nil { ds.Tunnels.Close() } + if ds.PeerSessions != nil { + ds.PeerSessions.Close() + } + if ds.PeerPool != nil { + ds.PeerPool.Close() + } // Then the membership, so the peers see this replica leave before its // services stop. ds.Membership.Stop() @@ -557,20 +578,31 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Create ModelRouterAdapter to wire into ModelLoader modelAdapter := nodes.NewModelRouterAdapter(router) - membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID) + // A replica proves who it is to its peers with a credential that it mints + // itself. Only the hash is published, in the same row as the address that + // the peers dial. + peerCred := cluster.NewPeerCredential() + membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID, peerAddress(cfg), peerCred) if err != nil { return nil, err } - tunnels := tunnel.NewRegistry(cluster.NewRegistry(authDB), cfg.Distributed.InstanceID) + clusterReg := cluster.NewRegistry(authDB) + tunnels := tunnel.NewRegistry(clusterReg, cfg.Distributed.InstanceID) // If a peer sweeps this replica while it stalls, the claims of the tunnels // that it still holds are written again when it registers again. membership.SetReclaimer(tunnels) + peerPool := tunnel.NewPeerPool(cfg.Distributed.InstanceID, peerCred, clusterReg) + peerSessions := tunnel.NewPeerSessions(tunnel.NewRelay(tunnels).Stream) success = true return &DistributedServices{ Membership: membership, Carriers: carrierStore, Tunnels: tunnels, + Instances: clusterReg, + PeerSessions: peerSessions, + PeerPool: peerPool, + WorkerDialer: tunnel.NewWorkerDialer(tunnels, peerPool), Broadcaster: broadcaster, WorkQueue: workQueue, AgentControl: carrier.NewAgents(active), @@ -600,7 +632,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // instances table and keeps its row fresh until the Membership is stopped. // Every replica does this on every carrier, because the change of carrier needs // to know which replicas are alive and ready. -func startMembership(ctx context.Context, db *gorm.DB, instanceID string) (*cluster.Membership, error) { +func startMembership(ctx context.Context, db *gorm.DB, instanceID, advertisedAddr string, peerCred cluster.PeerCredential) (*cluster.Membership, error) { if ctx == nil { ctx = context.Background() } @@ -610,6 +642,7 @@ func startMembership(ctx context.Context, db *gorm.DB, instanceID string) (*clus return nil, fmt.Errorf("migrating cluster tables: %w", err) } membership := cluster.NewMembership(cluster.NewRegistry(db), instanceID, internal.PrintableVersion()) + membership.SetPeer(advertisedAddr, peerCred) if err := membership.Start(ctx); err != nil { return nil, fmt.Errorf("registering this replica: %w", err) } @@ -619,3 +652,40 @@ func startMembership(ctx context.Context, db *gorm.DB, instanceID string) (*clus func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } + +// peerAddress returns the address at which the other replicas dial this one. An +// address that the operator configured is checked and used as it is. Without +// one, it comes from the route to the database, which all replicas share. +// +// A replica that has no usable address is not refused. It publishes none, and +// the replicas that need to relay through it find it unreachable, which is +// reported when they try. The log line says how to fix it. +func peerAddress(cfg *config.ApplicationConfig) string { + if configured := cfg.Distributed.PeerAddress; configured != "" { + reason, err := cluster.CheckAdvertisedAddr(configured) + if err != nil { + xlog.Warn("The configured peer address is unusable, so this replica publishes none and the other replicas cannot reach it", "address", configured, "error", err) + return "" + } + if reason != "" { + xlog.Warn("The configured peer address may not work for the other replicas", "address", configured, "reason", reason) + } + return configured + } + _, portText, err := net.SplitHostPort(cfg.APIAddress) + if err != nil { + xlog.Warn("The listen address has no port, so this replica publishes no peer address; set LOCALAI_PEER_ADDRESS", "address", cfg.APIAddress, "error", err) + return "" + } + port, err := strconv.Atoi(portText) + if err != nil { + xlog.Warn("The listen port is not a number, so this replica publishes no peer address; set LOCALAI_PEER_ADDRESS", "port", portText, "error", err) + return "" + } + addr, err := cluster.DiscoverAdvertisedAddr(cfg.Auth.DatabaseURL, port) + if err != nil { + xlog.Warn("This replica publishes no peer address, so a worker tunnel that it holds cannot be reached through the other replicas; set LOCALAI_PEER_ADDRESS", "error", err) + return "" + } + return addr +} diff --git a/core/cli/run.go b/core/cli/run.go index 197f638500a4..e64cf2321865 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -165,6 +165,7 @@ type RunCMD struct { // Distributed / Horizontal Scaling Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"` InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"` + PeerAddress string `env:"LOCALAI_PEER_ADDRESS" help:"host:port at which the other frontends dial this one to reach a worker tunnel it holds (found from the route to the database if empty)" group:"distributed"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` @@ -363,6 +364,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.InstanceID != "" { opts = append(opts, config.WithDistributedInstanceID(r.InstanceID)) } + if r.PeerAddress != "" { + opts = append(opts, config.WithDistributedPeerAddress(r.PeerAddress)) + } if r.NatsURL != "" { opts = append(opts, config.WithNatsURL(r.NatsURL)) } diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index a6e11ffdec52..b02d666a095b 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -13,8 +13,12 @@ import ( // DistributedConfig holds configuration for horizontal scaling mode. // When Enabled is true, PostgreSQL and NATS are required. type DistributedConfig struct { - Enabled bool // --distributed / LOCALAI_DISTRIBUTED - InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + Enabled bool // --distributed / LOCALAI_DISTRIBUTED + InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + // PeerAddress is the host and port at which the other frontends dial this + // one, to reach a worker tunnel that this one holds. Empty means that the + // address is found from the route to the database. LOCALAI_PEER_ADDRESS. + PeerAddress string NatsURL string // --nats-url / LOCALAI_NATS_URL StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) @@ -196,6 +200,14 @@ func WithDistributedInstanceID(id string) AppOption { } } +// WithDistributedPeerAddress sets the address at which other frontends dial this +// one. +func WithDistributedPeerAddress(addr string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.PeerAddress = addr + } +} + func WithNatsURL(url string) AppOption { return func(o *ApplicationConfig) { o.Distributed.NatsURL = url diff --git a/core/http/app.go b/core/http/app.go index 40eb11766d68..00362ce39d75 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -28,6 +28,7 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/finetune" "github.com/mudler/LocalAI/core/services/galleryop" @@ -579,7 +580,15 @@ func API(application *application.Application) (*echo.Echo, error) { } natsCfg := distCfg.NatsAuthConfig() routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg, registerOpts...) - routes.RegisterClusterRoutes(e, registry, tunnels) + var ( + instances *cluster.Registry + onPeer func(string, *tunnel.Session) + ) + if d := application.Distributed(); d != nil { + instances = d.Instances + onPeer = d.PeerSessions.Accept + } + routes.RegisterClusterRoutes(e, registry, tunnels, instances, onPeer) routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDial, registerOpts...) // Distributed SSE routes (job progress + agent events via NATS) diff --git a/core/http/auth/helpers_test.go b/core/http/auth/helpers_test.go index 289b3692d6ee..8490211f9111 100644 --- a/core/http/auth/helpers_test.go +++ b/core/http/auth/helpers_test.go @@ -128,7 +128,7 @@ func newNodeSelfServiceTestApp(db *gorm.DB, appConfig *config.ApplicationConfig, func newClusterTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { e := echo.New() e.Use(auth.Middleware(db, appConfig)) - routes.RegisterClusterRoutes(e, nil, nil) + routes.RegisterClusterRoutes(e, nil, nil, nil, nil) e.GET("/api/cluster/other", ok) return e } diff --git a/core/http/auth/middleware_test.go b/core/http/auth/middleware_test.go index 0fa7ab4bd507..6f6e729c6e77 100644 --- a/core/http/auth/middleware_test.go +++ b/core/http/auth/middleware_test.go @@ -456,6 +456,23 @@ var _ = Describe("Auth Middleware", func() { Expect(other.Code).To(Equal(http.StatusUnauthorized)) }) + It("delegates the peer link to its own credential check, and only that path", func() { + app = newClusterTestApp(db, appConfig) + + // The header of the peer credential, which is no API key and no + // session. The middleware lets the dial through, and the handler (which + // has no table of replicas here) answers 503. If the middleware stopped + // it, the answer would be 401. + peer := doRequest(app, http.MethodGet, "/api/cluster/peer?id=replica-a", func(r *http.Request) { + r.Header.Set("X-LocalAI-Peer-Token", "own-peer-credential") + }) + Expect(peer.Code).To(Equal(http.StatusServiceUnavailable)) + + // A dial with no credential is refused whoever answers. + anonymous := doRequest(app, http.MethodGet, "/api/cluster/peer?id=replica-a") + Expect(anonymous.Code).To(Equal(http.StatusUnauthorized)) + }) + It("lets downstream node auth reject missing and invalid registration tokens", func() { const registrationToken = "node-registration-secret" app = newNodeSelfServiceTestApp(db, appConfig, registrationToken) diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index fcef6e0d2281..55ea8d532da9 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -91,8 +91,13 @@ func isPublicRoute(method, path string) bool { // it with tunnel.ConnectPath. const ClusterConnectPath = "/api/cluster/connect" +// ClusterPeerPath is the route that a replica dials to reach a tunnel that +// another replica holds. The handler checks the own credential of the dialling +// replica. The spec of the peer endpoint compares it with tunnel.PeerPath. +const ClusterPeerPath = "/api/cluster/peer" + // usesAlternativeAuthentication identifies requests whose credentials are // validated by route-group middleware instead of the global auth middleware. func usesAlternativeAuthentication(path string) bool { - return strings.HasPrefix(path, "/api/node/") || path == ClusterConnectPath + return strings.HasPrefix(path, "/api/node/") || path == ClusterConnectPath || path == ClusterPeerPath } diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go new file mode 100644 index 000000000000..8dd8a55b3fef --- /dev/null +++ b/core/http/endpoints/cluster/peer.go @@ -0,0 +1,122 @@ +package cluster + +import ( + "errors" + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/xlog" + + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/tunnel" +) + +// PeerHandler serves the link that a replica dials to reach a worker tunnel that +// another replica holds. The dial becomes a websocket, the websocket becomes a +// multiplexed session, and onSession gets the session. +// +// One credential decides, and it is the credential of the dialling replica. The +// id in ?id= is a claim. The handler resolves it to a row of the instances table +// and checks the secret in the header against the hash that the row publishes. +// The registration token of the deployment opens no peer link: it is held by +// every worker, and a worker that could pose as a replica could relay to every +// tunnel that a replica owns. +// +// All checks come before the upgrade, and the credential comes first. A dial with +// no credential is the anonymous case and gets 401 whatever else is missing. +// +// onSession runs on the goroutine of the request, so it returns promptly. The +// session outlives the handler, because the upgrade hijacks the connection, and +// onSession owns closing it. +func PeerHandler(instances *clustersvc.Registry, onSession func(peerID string, sess *tunnel.Session)) echo.HandlerFunc { + // The default origin check of gorilla refuses a browser of another origin and + // admits a client with no Origin header, which every peer is. + upgrader := tunnel.NewUpgrader() + + return func(c echo.Context) error { + presented := c.Request().Header.Get(clustersvc.PeerIdentityHeader) + if presented == "" { + // A missing header is refused and not waved through. Accepting a dial + // with no credential would leave the hole as open as it was, because an + // attacker omits the header too. + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // Not 401. A frontend with no registry cannot resolve any replica, and + // "unauthorized" would send an operator after a problem with a credential + // that does not exist. The route is registered in every deployment so that + // the route coverage test sees it, and a handler that treated a nil registry + // as "nothing to check" would publish an unauthenticated multiplexer. + if instances == nil { + return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled") + } + + peerID := c.QueryParam("id") + if peerID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing peer id") + } + + inst, err := instances.Get(c.Request().Context(), peerID) + switch { + case errors.Is(err, clustersvc.ErrInstanceNotFound): + // 401 and not 404, so that a caller cannot list replica ids by status. + xlog.Warn("Refusing a peer link: the dial named a replica this deployment has no row for", "peer", peerID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + case err != nil: + // A query that failed is neither a rejection nor an absence. 401 on an + // unreadable database would tell a healthy replica that its credential + // is wrong. + xlog.Error("Looking up a replica for its peer dial failed", "peer", peerID, "error", err) + return echo.NewHTTPError(http.StatusInternalServerError, "peer lookup failed") + } + + // Two operator problems with two cures. An empty stored hash is a replica + // that has not registered a credential, so there is nothing to check and it + // must register again, which a restart does. A mismatch is a dialler that + // presents the wrong secret. An empty hash never matches: it is not "no + // restriction". + if inst.PeerTokenHash == "" { + xlog.Warn("Refusing a peer link: the replica it claims to be has no peer credential published, so nothing here can verify the claim", + "peer", peerID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + if !clustersvc.PeerTokenMatches(presented, inst.PeerTokenHash) { + xlog.Warn("Refusing a peer link: the dial presented the wrong credential for the replica it claims to be", "peer", peerID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has written its own failure to the client. + xlog.Debug("tunnel peer link upgrade failed", "peer", peerID, "error", err) + return nil + } + + sess, err := tunnel.PeerServerSession(ws) + if err != nil { + xlog.Error("Tunnel peer link session setup failed", "peer", peerID, "error", err) + _ = ws.Close() + return nil + } + + if onSession == nil { + // Nothing reads this session, so the peer must not believe it has a + // live link. + _ = sess.Close() + return nil + } + + xlog.Debug("tunnel peer link established", "peer", peerID, "remote", ws.RemoteAddr().String()) + // net/http recovers a panic of this goroutine but does not close the + // hijacked connection, so a panicking callback would leave the peer with a + // link that nobody accepts streams on. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() + onSession(peerID, sess) + return nil + } +} diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go new file mode 100644 index 000000000000..d2d5ec58e63d --- /dev/null +++ b/core/http/endpoints/cluster/peer_test.go @@ -0,0 +1,305 @@ +package cluster_test + +import ( + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "time" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/auth" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// These specs drive the production handler over a real websocket and a real +// yamux session, because what they pin is a property of the place where the +// connection is hijacked: the id of a peer is a query parameter, and whether it +// is checked is decided there. +var _ = Describe("Peer link endpoint", func() { + const ( + ownerID = "replica-owner" + callerID = "replica-caller" + // The registration token is held by every worker. Holding it must open + // no peer link. + registrationToken = "shared-registration-token" + ) + + var ( + ctx context.Context + db *gorm.DB + reg *cluster.Registry + sessions *tunnel.PeerSessions + srv *httptest.Server + ownerCred cluster.PeerCredential + callerCred cluster.PeerCredential + ) + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + + // The owner holds a worker, and accepts peer links. + ownerTunnels := tunnel.NewRegistry(reg, ownerID) + sessions = tunnel.NewPeerSessions(tunnel.NewRelay(ownerTunnels).Stream) + DeferCleanup(sessions.Close) + + e := echo.New() + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(reg, sessions.Accept)) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + + ownerCred, callerCred = cluster.NewPeerCredential(), cluster.NewPeerCredential() + addr := strings.TrimPrefix(srv.URL, "http://") + Expect(reg.RegisterPeer(ctx, ownerID, "test", 0, "", addr, ownerCred.Hash())).To(Succeed()) + Expect(reg.RegisterPeer(ctx, callerID, "test", 0, "", "127.0.0.1:9", callerCred.Hash())).To(Succeed()) + }) + + // raw dials the route with the headers a spec chooses. + raw := func(query string, header http.Header) (*http.Response, error) { + GinkgoHelper() + url := "ws" + strings.TrimPrefix(srv.URL, "http") + tunnel.PeerPath + query + ws, resp, err := tunnel.NewDialer(5*time.Second).Dial(url, header) + if ws != nil { + DeferCleanup(func() { _ = ws.Close() }) + } + return resp, err + } + + Describe("the answers before the upgrade", func() { + It("answers 401 to a dial with no credential, whatever else is missing", func() { + resp, err := raw("", nil) + Expect(err).To(MatchError(websocket.ErrBadHandshake)) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 401 to a dial that holds only the registration token", func() { + header := http.Header{"Authorization": []string{"Bearer " + registrationToken}} + resp, err := raw("?id="+callerID, header) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 503 to a credential when the frontend has no table of replicas", func() { + e := echo.New() + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(nil, nil)) + bare := httptest.NewServer(e) + DeferCleanup(bare.Close) + req, _ := http.NewRequest(http.MethodGet, bare.URL+tunnel.PeerPath+"?id=x", nil) + req.Header.Set(cluster.PeerIdentityHeader, "anything") + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + }) + + It("answers 400 to a credential with no replica id", func() { + resp, err := raw("", http.Header{cluster.PeerIdentityHeader: []string{callerCred.Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("answers 401, not 404, for a replica it has no row for", func() { + resp, err := raw("?id=ghost", http.Header{cluster.PeerIdentityHeader: []string{callerCred.Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 401 for a replica whose row publishes no credential", func() { + Expect(db.Model(&cluster.Instance{}).Where("id = ?", callerID).Update("peer_token_hash", "").Error).To(Succeed()) + resp, err := raw("?id="+callerID, http.Header{cluster.PeerIdentityHeader: []string{callerCred.Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), "an empty hash authorises nobody") + }) + + It("answers 401 for the id of a replica with the credential of another", func() { + resp, err := raw("?id="+callerID, http.Header{cluster.PeerIdentityHeader: []string{cluster.NewPeerCredential().Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(sessions.Holds(callerID)).To(BeFalse()) + }) + }) + + Describe("a replica that proves its id", func() { + It("gets a link, and the owner holds exactly one session for it", func() { + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + st, err := pool.Open(ctx, ownerID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + Eventually(func() bool { return sessions.Holds(callerID) }, "10s").Should(BeTrue()) + }) + + It("is refused when it declares another replica's id, and cannot evict the link of the real one", func() { + good := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(good.Close) + st, err := good.Open(ctx, ownerID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + Eventually(func() bool { return sessions.Holds(callerID) }, "10s").Should(BeTrue()) + + impostor := tunnel.NewPeerPool(callerID, cluster.NewPeerCredential(), reg) + DeferCleanup(impostor.Close) + _, err = impostor.Open(ctx, ownerID) + Expect(err).To(MatchError(tunnel.ErrPeerRejected)) + Expect(err).To(MatchError(tunnel.ErrPeerUnreachable), "every caller that handles an unreachable peer handles a refusal") + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse(), "a refusal is not absence") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse()) + Expect(sessions.Holds(callerID)).To(BeTrue(), "the link of the real replica stays") + }) + + It("reaches a worker that the owner holds, through the link, with the bytes intact", func() { + // A worker behind the owner. + ownerTunnels := tunnel.NewRegistry(reg, ownerID) + worker := &workerEcho{} + sessions.Close() + sessions = tunnel.NewPeerSessions(tunnel.NewRelay(ownerTunnels).Stream) + DeferCleanup(sessions.Close) + e := echo.New() + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(reg, sessions.Accept)) + srv2 := httptest.NewServer(e) + DeferCleanup(srv2.Close) + Expect(reg.RegisterPeer(ctx, ownerID, "test", 0, "", strings.TrimPrefix(srv2.URL, "http://"), ownerCred.Hash())).To(Succeed()) + + front, back := workerSessions() + worker.serve(back) + _, err := ownerTunnels.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + + callerTunnels := tunnel.NewRegistry(reg, callerID) + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + dialer := tunnel.NewWorkerDialer(callerTunnels, pool) + + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + big := strings.Repeat("0123456789abcdef", 1<<14) // 256 KiB + go func() { _, _ = conn.Write([]byte(big)) }() + got := make([]byte, len(big)) + Expect(conn.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = io.ReadFull(conn, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal(big)) + }) + + It("reports a peer with no listener as unreachable, and a peer without an address as unreachable too", func() { + Expect(reg.RegisterPeer(ctx, "silent", "test", 0, "", "127.0.0.1:1", cluster.NewPeerCredential().Hash())).To(Succeed()) + Expect(reg.Register(ctx, "addressless", "test", 0, "")).To(Succeed()) + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + for _, peer := range []string{"silent", "addressless"} { + _, err := pool.Open(ctx, peer) + Expect(err).To(MatchError(tunnel.ErrPeerUnreachable), peer) + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse(), peer) + } + _, err := pool.Open(ctx, "no-such-replica") + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeFalse(), "a replica that is not registered is not an unreachable one") + }) + + It("reports a closed pool as itself", func() { + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + pool.Close() + pool.Close() + _, err := pool.Open(ctx, ownerID) + Expect(err).To(MatchError(tunnel.ErrPoolClosed)) + }) + + It("does not blame the peer when the budget of the caller runs out", func() { + // A listener that accepts the TCP connection and never answers the upgrade. + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = ln.Close() }) + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + DeferCleanup(func() { _ = c.Close() }) + } + }() + Expect(reg.RegisterPeer(ctx, "slow", "test", 0, "", ln.Addr().String(), cluster.NewPeerCredential().Hash())).To(Succeed()) + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + short, cancel := context.WithTimeout(ctx, 200*time.Millisecond) + defer cancel() + _, err = pool.Open(short, "slow") + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeFalse()) + }) + }) +}) + +var _ = Describe("The path of the peer route", func() { + It("is the one that the global auth middleware lets through", func() { + Expect(auth.ClusterPeerPath).To(Equal(tunnel.PeerPath)) + }) +}) + +// workerEcho is a worker that accepts any stream and echoes. +type workerEcho struct{} + +func (workerEcho) serve(sess *tunnel.Session) { + go func() { + for { + st, err := sess.AcceptStream() + if err != nil { + return + } + go func() { + defer func() { _ = st.Close() }() + if _, _, err := tunnel.ReadStreamRequest(st); err != nil { + return + } + if err := tunnel.WriteStreamAccepted(st); err != nil { + return + } + _, _ = io.Copy(st, st) + }() + } + }() +} + +// workerSessions opens a real websocket on a loopback socket and returns the +// frontend side and the worker side of an inference session on it. +func workerSessions() (front, back *tunnel.Session) { + GinkgoHelper() + up := tunnel.NewUpgrader() + got := make(chan *tunnel.Session, 1) + s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := tunnel.ServerSession(ws, tunnel.LaneInference) + if err != nil { + _ = ws.Close() + return + } + got <- sess + })) + DeferCleanup(s.Close) + ws, _, err := tunnel.NewDialer(5*time.Second).Dial("ws"+strings.TrimPrefix(s.URL, "http"), nil) + Expect(err).ToNot(HaveOccurred()) + back, err = tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = back.Close() }) + Eventually(got).Should(Receive(&front)) + DeferCleanup(func() { _ = front.Close() }) + return front, back +} diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go index 5e1df77513b2..34bc1484721e 100644 --- a/core/http/routes/cluster.go +++ b/core/http/routes/cluster.go @@ -3,6 +3,7 @@ package routes import ( "github.com/labstack/echo/v4" clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/tunnel" ) @@ -15,6 +16,12 @@ import ( // refuses every worker that has no tunnel credential. A frontend that is not // distributed answers 503 to a caller that has a credential and 401 to one that // has none. -func RegisterClusterRoutes(e *echo.Echo, registry *nodes.NodeRegistry, tunnels *tunnel.Registry) { +// +// The peer route is registered in the same way, for the same reason: the route +// coverage test walks it, and it answers 401 to a dial with no credential. It +// checks the credential of the dialling replica against the instances table, and +// a deployment with no table answers 503. +func RegisterClusterRoutes(e *echo.Echo, registry *nodes.NodeRegistry, tunnels *tunnel.Registry, instances *cluster.Registry, onPeer func(string, *tunnel.Session)) { e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(registry, tunnels)) + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(instances, onPeer)) } diff --git a/core/services/nodes/control_tunnel_test.go b/core/services/nodes/control_tunnel_test.go index 18811b257a70..4fa19cf88793 100644 --- a/core/services/nodes/control_tunnel_test.go +++ b/core/services/nodes/control_tunnel_test.go @@ -135,8 +135,12 @@ var _ = Describe("The tunnel carrier of the control verbs", func() { worker.replyJSON(workerctl.VerbBackendStop, workerctl.BackendStopReply{Success: true}) worker.replyJSON(workerctl.VerbBackendDelete, workerctl.BackendDeleteReply{Success: true}) worker.replyJSON(workerctl.VerbModelDelete, workerctl.ModelDeleteReply{Success: true}) - worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") }) - worker.on(workerctl.VerbBackendUpgrade, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") }) + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") + }) + worker.on(workerctl.VerbBackendUpgrade, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"reply":{"success":true}}`+"\n") + }) _, err := control.StopLoadOperation(context.Background(), node, workerctl.ModelStopRequest{ProcessKey: "m#0", OperationID: "op"}) Expect(err).ToNot(HaveOccurred()) From 67a417a056f6907c08c30e52c1433d8cde8b4603 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:12:40 +0000 Subject: [PATCH 060/162] fix(tunnel): report a bulk lane that is down as that through a relay A worker whose bulk session is down still carries model calls on its inference lane. A replica that asked the owner for the bulk lane only got "the owner could not open a stream", which reads as a route that does not exist, and a scheduler that demoted the worker for it would take away the requests that work. The owner now says that it is the bulk lane, with its own refusal code, and the dialling replica reports ErrNoBulkSession without ErrNoRoute. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/dialer_test.go | 7 ++++--- core/services/tunnel/relay.go | 19 ++++++++++++++++--- core/services/tunnel/relay_test.go | 1 + 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/core/services/tunnel/dialer_test.go b/core/services/tunnel/dialer_test.go index d54cf301ab4b..751ad7a708a1 100644 --- a/core/services/tunnel/dialer_test.go +++ b/core/services/tunnel/dialer_test.go @@ -248,7 +248,7 @@ var _ = Describe("The worker dialer", func() { Expect(peers.requested()).To(ContainElement(ContainSubstring("bulk-only"))) }) - It("gets the refusal of the owner when the owner has no bulk session, and not a transfer on the other lane", func() { + It("gets the refusal of the owner when the owner has no bulk session, as that and not as a missing route", func() { Expect(other.Disconnect("w1")).To(BeTrue()) front, back := sessionPair(tunnel.LaneInference) worker.serve(back) @@ -258,8 +258,9 @@ var _ = Describe("The worker dialer", func() { dialer := tunnel.NewWorkerDialer(tunnels, peers) _, err = dialer.Dial(ctx, "w1", tunnel.StreamTagHTTP, "", tunnel.WithBulkLane()) Expect(err).To(HaveOccurred()) - Expect(errors.Is(err, tunnel.ErrRelayUnavailable)).To(BeTrue(), "%v", err) - Expect(worker.seen()).To(BeEmpty()) + Expect(errors.Is(err, tunnel.ErrNoBulkSession)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "the tunnel is held and carries model calls; only a lane is down") + Expect(worker.seen()).To(BeEmpty(), "nothing may have crossed the inference lane") }) It("states the remaining budget of the caller in the relay request", func() { diff --git a/core/services/tunnel/relay.go b/core/services/tunnel/relay.go index dd44228a9e09..dd8df8d6e42a 100644 --- a/core/services/tunnel/relay.go +++ b/core/services/tunnel/relay.go @@ -35,6 +35,7 @@ const ( relayReplyAccepted = "relay-ok" relayCodeNotOwner = "relay-not-owner" relayCodeUnavailable = "relay-unavailable" + relayCodeNoBulk = "relay-no-bulk" relayCodeBadRequest = "relay-bad-request" ) @@ -65,6 +66,12 @@ func relayLane(o dialOptions) string { // carry a stream, so a retry is worth something and looking elsewhere is not. // ErrRelayRequestInvalid is a bug of the caller, and no retry helps. // +// The owner also reports a bulk lane that is down (ErrNoBulkSession), so that the +// dialling replica reports it as that and not as a route that does not exist. A +// worker whose bulk session is being dialled again has a tunnel that carries +// model calls, and a scheduler that demoted it for the lane would lose it for the +// requests that work. +// // None of them is built over an absence error. A refusal proves that a replica // answered. var ( @@ -150,6 +157,8 @@ func WriteRelayRefusal(w io.Writer, reason error) error { code = relayCodeNotOwner case errors.Is(reason, ErrRelayUnavailable): code = relayCodeUnavailable + case errors.Is(reason, ErrNoBulkSession): + code = relayCodeNoBulk } text := "" if reason != nil { @@ -191,6 +200,8 @@ func ReadRelayReply(r io.Reader) error { return fmt.Errorf("%w: %s", ErrNotOwner, text) case relayCodeUnavailable: return fmt.Errorf("%w: %s", ErrRelayUnavailable, text) + case relayCodeNoBulk: + return fmt.Errorf("%w: %s", ErrNoBulkSession, text) case relayCodeBadRequest: return fmt.Errorf("%w: %s", ErrRelayRequestInvalid, text) } @@ -324,9 +335,11 @@ func (r *Relay) accept(peerID string, stream net.Conn) (net.Conn, bool) { } local, err := r.tunnels.Open(ctx, nodeID, openLane, openOpts...) if err != nil { - if errors.Is(err, ErrNotOwner) { - // As itself: the worker is probably well on another replica, and this - // is the answer that tells the caller to look there. + if errors.Is(err, ErrNotOwner) || errors.Is(err, ErrNoBulkSession) { + // As itself. ErrNotOwner: the worker is probably well on another + // replica, and this is the answer that tells the caller to look there. + // ErrNoBulkSession: the tunnel is held here and only its bulk lane is + // down, which the caller must not read as a failure of this replica. r.refuse(peerID, stream, err) return nil, false } diff --git a/core/services/tunnel/relay_test.go b/core/services/tunnel/relay_test.go index 3c23efa5d6b8..903f2c8acbd9 100644 --- a/core/services/tunnel/relay_test.go +++ b/core/services/tunnel/relay_test.go @@ -97,6 +97,7 @@ var _ = Describe("The relay reply", func() { It("keeps the three refusals apart", func() { Expect(roundTrip(tunnel.ErrNotOwner)).To(MatchError(tunnel.ErrNotOwner)) Expect(roundTrip(tunnel.ErrRelayUnavailable)).To(MatchError(tunnel.ErrRelayUnavailable)) + Expect(roundTrip(tunnel.ErrNoBulkSession)).To(MatchError(tunnel.ErrNoBulkSession)) Expect(roundTrip(errors.New("anything else"))).To(MatchError(tunnel.ErrRelayRequestInvalid)) }) From d94aa415a8848e2955ee7cce397b3ff80550aa47 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:12:40 +0000 Subject: [PATCH 061/162] feat(grpc): let a client reach its backend through a dialer and say why a dial failed A client can take its own dialer, so that a backend of a worker that holds a tunnel is reached by a stream on that tunnel and not by a TCP connection to an address. The passthrough resolver keeps gRPC from resolving a name that is not a host. gRPC flattens a failure of the dialer into codes.Unavailable, the code of a backend that died. The client keeps the error of its last dial, and TransportFailureOf tells a caller that decides whether a backend is dead whether the dial failed in the transport or the host answered about the backend. LastDialErrorOf looks through decorators that implement Unwrap. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- pkg/grpc/backend.go | 100 +++++++++++++++++++++++++++++- pkg/grpc/client.go | 61 +++++++++++++++++- pkg/grpc/dialer_test.go | 133 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 292 insertions(+), 2 deletions(-) create mode 100644 pkg/grpc/dialer_test.go diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index ac5b95f28efe..9f52f0ea3087 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -2,6 +2,8 @@ package grpc import ( "context" + "errors" + "net" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "google.golang.org/grpc" @@ -29,7 +31,103 @@ func NewClientWithToken(address string, parallel bool, wd WatchDog, enableWatchD return buildClient(address, parallel, wd, enableWatchDog, token) } -func buildClient(address string, parallel bool, wd WatchDog, enableWatchDog bool, token string) Backend { +// NewClientWithDialer creates a gRPC client that reaches its backend through +// dialer and not through a TCP connection to address. The tunnel carrier uses +// it, and the address is then the name of a backend process of a worker. +// +// The outcome of every dial is recorded, so that a caller can ask why a call +// failed: see LastDialErrorOf. +func NewClientWithDialer(address string, parallel bool, wd WatchDog, enableWatchDog bool, token string, dialer func(ctx context.Context, addr string) (net.Conn, error)) Backend { + if bc, ok := embeds[address]; ok { + return bc + } + // Assigned on the concrete type and not through an assertion that can fail: + // a failed assertion would hand back a client that dials the address + // directly, which is the bypass this constructor exists to close. + c := buildClient(address, parallel, wd, enableWatchDog, token) + c.dialer = func(ctx context.Context, addr string) (net.Conn, error) { + conn, err := dialer(ctx, addr) + c.recordDialErr(err) + return conn, err + } + return c +} + +// DialErrorReporter is implemented by a Backend that reaches its process through +// a custom transport and can say whether the transport failed. It is not part of +// Backend: the few callers that act on the difference ask for it, and widening +// Backend would make every wrapper and every double implement a method that they +// have no answer for. +type DialErrorReporter interface { + LastDialError() error +} + +// BackendAnswer is implemented by the error of a custom dialer when the dial +// reached the host of the backend and the host answered. A refusal of the worker +// to open a stream to a backend that is gone is one: it is evidence about the +// backend, like a refused connection on a direct dial. An error that does not +// implement it is a failure of the transport, and nothing was learned about the +// backend. +type BackendAnswer interface { + error + IsBackendAnswer() bool +} + +// BackendUnwrapper is implemented by a Backend that decorates another one. +// Every decorator must implement it: a decorator that embeds the Backend +// interface inherits what Backend declares and nothing else, and +// DialErrorReporter is not declared there. A wrapped client would then stop +// answering whether the transport failed, and the guards built on the answer +// would read nil in production while every spec that built a raw client passed. +type BackendUnwrapper interface { + Unwrap() Backend +} + +// maxBackendUnwrapDepth bounds the walk of LastDialErrorOf. It guards against a +// cycle that a future wrapper could introduce, and nothing real nests this deep. +const maxBackendUnwrapDepth = 16 + +// LastDialErrorOf reports why the most recent dial under b failed, looking +// through any decorators, or nil when the dial succeeded or nothing under b has a +// custom transport. +func LastDialErrorOf(b Backend) error { + for range maxBackendUnwrapDepth { + if b == nil { + return nil + } + if reporter, ok := b.(DialErrorReporter); ok { + return reporter.LastDialError() + } + wrapper, ok := b.(BackendUnwrapper) + if !ok { + return nil + } + b = wrapper.Unwrap() + } + return nil +} + +// TransportFailureOf is LastDialErrorOf for a caller that decides whether a +// backend is dead. It returns the error of the last dial when that error is a +// failure of the transport, and nil when the dial succeeded or when the host +// answered about the backend (see BackendAnswer). +// +// A caller that gets an error from it must not conclude anything about the +// backend: not that it is dead, and not that it is well. It must not reap the +// row of a model, evict a cached model or count a miss. +func TransportFailureOf(b Backend) error { + err := LastDialErrorOf(b) + if err == nil { + return nil + } + var answer BackendAnswer + if errors.As(err, &answer) && answer.IsBackendAnswer() { + return nil + } + return err +} + +func buildClient(address string, parallel bool, wd WatchDog, enableWatchDog bool, token string) *Client { if !enableWatchDog { wd = nil } diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index f46b9dcf03a7..e48f60cd43cd 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "net" "sync" "time" @@ -32,6 +33,20 @@ type Client struct { inFlight int parallel bool token string + // dialer replaces the transport that gRPC would use to reach address. In a + // deployment where workers hold a tunnel, it opens a stream on that tunnel, + // and address stops being a socket and becomes the name of a backend process + // inside the worker. Nil keeps the TCP dial of gRPC, which is what every + // caller outside a tunnel wants. + dialer func(ctx context.Context, addr string) (net.Conn, error) + + // dialErrMu guards lastDialErr. It is not the embedded mutex: that one + // guards inFlight and is taken on every call, and a dialer runs under the + // machinery of gRPC, where taking it again is not something this type can + // reason about. + dialErrMu sync.Mutex + lastDialErr error + sync.Mutex opMutex sync.Mutex wd WatchDog @@ -80,7 +95,17 @@ func (c *Client) dial() (*grpc.ClientConn, error) { if c.token != "" { opts = append(opts, grpc.WithPerRPCCredentials(bearerToken{token: c.token})) } - return grpc.NewClient(c.address, opts...) + target := c.address + if c.dialer != nil { + // The address still goes to grpc.NewClient, because it names the target + // in every error message and in the authority header. What it no longer + // decides is where the bytes go. The passthrough resolver keeps gRPC from + // resolving it: the name of a backend process of a worker is not a host + // that a resolver knows. + opts = append(opts, grpc.WithContextDialer(c.dialer)) + target = "passthrough:///" + c.address + } + return grpc.NewClient(target, opts...) } func (c *Client) HealthCheck(ctx context.Context) (bool, error) { @@ -1424,3 +1449,37 @@ func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts .. client := pb.NewBackendClient(conn) return client.ModelMetadata(ctx, in, opts...) } + +// LastDialError returns the error of the most recent attempt of the custom +// dialer of this client, or nil when that attempt succeeded or there is no custom +// dialer. +// +// It exists because gRPC destroys the difference that its callers need. A failure +// of the dialer reaches a call as codes.Unavailable with the cause flattened into +// a message, and codes.Unavailable is also what a backend process that died +// produces. The two call for opposite actions: the row of a dead backend is +// reaped, and a transport that could not reach a live backend must never cause +// that. Recording the error here keeps the error value for the caller. +// +// It is the last dial on this client and not the last dial for one call. The +// callers that build a client for one probe and close it get an exact answer. +// The one that reads the long-lived client of a model reads it after the call +// failed, so a concurrent dial can record or clear the value in between. A dial +// that failed in that window makes a dead backend look unreachable for now, and +// its row survives one more round. A dial that succeeded clears the value, and +// a failure of the transport reads as a failure of the backend, which is what +// happened before the dialer existed. Neither is a new hazard. +func (c *Client) LastDialError() error { + c.dialErrMu.Lock() + defer c.dialErrMu.Unlock() + return c.lastDialErr +} + +// recordDialErr stores the outcome of one dial. A success clears the failure +// before it, so that a client that recovered does not keep reporting an error +// that describes nothing. +func (c *Client) recordDialErr(err error) { + c.dialErrMu.Lock() + c.lastDialErr = err + c.dialErrMu.Unlock() +} diff --git a/pkg/grpc/dialer_test.go b/pkg/grpc/dialer_test.go new file mode 100644 index 000000000000..74e88be55a31 --- /dev/null +++ b/pkg/grpc/dialer_test.go @@ -0,0 +1,133 @@ +package grpc + +import ( + "context" + "errors" + "net" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// The dialer replaces the transport, and the client records why a dial failed. +// gRPC flattens a failure of the dialer into codes.Unavailable, the code of a +// backend that died, and the caller needs the difference to know whether to +// reap a row. +var _ = Describe("A client with its own dialer", func() { + var addr string + + BeforeEach(func() { + addr, _ = startTestServer("") + }) + + It("reaches the backend through the dialer and not through the address it is given", func() { + var dials atomic.Int32 + client := NewClientWithDialer("a-name-that-resolves-nowhere:1", false, nil, false, "", func(ctx context.Context, _ string) (net.Conn, error) { + dials.Add(1) + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + }) + + ok, err := client.HealthCheck(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(dials.Load()).To(BeNumerically(">=", 1)) + Expect(LastDialErrorOf(client)).To(BeNil()) + }) + + It("passes the address it was given to the dialer, as the name of the backend", func() { + var got atomic.Value + client := NewClientWithDialer("10.0.0.5:50051", false, nil, false, "", func(ctx context.Context, a string) (net.Conn, error) { + got.Store(a) + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + }) + _, err := client.HealthCheck(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Load()).To(Equal("10.0.0.5:50051")) + }) + + It("records the error of a dial that failed, with its value intact", func() { + boom := errors.New("no route") + client := NewClientWithDialer("10.0.0.5:50051", false, nil, false, "", func(context.Context, string) (net.Conn, error) { + return nil, boom + }) + ok, err := client.HealthCheck(context.Background()) + Expect(ok).To(BeFalse()) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(LastDialErrorOf(client), boom)).To(BeTrue(), "gRPC flattens the error; the client keeps the value") + Expect(errors.Is(TransportFailureOf(client), boom)).To(BeTrue()) + }) + + It("clears the error when the next dial succeeds", func() { + var fail atomic.Bool + fail.Store(true) + client := NewClientWithDialer("x:1", false, nil, false, "", func(ctx context.Context, _ string) (net.Conn, error) { + if fail.Load() { + return nil, errors.New("down") + } + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + }) + _, _ = client.HealthCheck(context.Background()) + Expect(LastDialErrorOf(client)).ToNot(BeNil()) + + fail.Store(false) + ok, err := client.HealthCheck(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(LastDialErrorOf(client)).To(BeNil()) + }) + + It("does not call the answer of a host a failure of the transport", func() { + refusal := &answerErr{errors.New("the worker could not reach the backend")} + client := NewClientWithDialer("x:1", false, nil, false, "", func(context.Context, string) (net.Conn, error) { + return nil, refusal + }) + _, _ = client.HealthCheck(context.Background()) + Expect(LastDialErrorOf(client)).To(MatchError(refusal), "the error is still recorded") + Expect(TransportFailureOf(client)).To(BeNil(), "but it is evidence about the backend, like a refused connection") + }) + + It("sees the dial through a decorator that unwraps", func() { + boom := errors.New("no route") + inner := NewClientWithDialer("x:1", false, nil, false, "", func(context.Context, string) (net.Conn, error) { return nil, boom }) + _, _ = inner.HealthCheck(context.Background()) + outer := &decorator{Backend: inner} + Expect(errors.Is(LastDialErrorOf(outer), boom)).To(BeTrue()) + }) + + It("answers nil for a decorator that does not unwrap, and for a client with no dialer", func() { + Expect(LastDialErrorOf(&undecorated{Backend: NewClient(addr, false, nil, false)})).To(BeNil()) + Expect(LastDialErrorOf(NewClient(addr, false, nil, false))).To(BeNil()) + Expect(LastDialErrorOf(nil)).To(BeNil()) + }) + + It("stops walking a cycle of decorators", func() { + a := &decorator{} + a.Backend = a + Expect(LastDialErrorOf(a)).To(BeNil()) + }) + + It("keeps the bearer token on every call", func() { + tokenAddr, _ := startTestServer("secret") + client := NewClientWithDialer("x:1", false, nil, false, "secret", func(ctx context.Context, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", tokenAddr) + }) + ok, err := client.HealthCheck(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + }) +}) + +type answerErr struct{ error } + +func (*answerErr) IsBackendAnswer() bool { return true } + +type decorator struct{ Backend } + +func (d *decorator) Unwrap() Backend { return d.Backend } + +type undecorated struct{ Backend } From c5fd85b5d6a10213767e747aac113af60a678181 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:12:40 +0000 Subject: [PATCH 062/162] fix(nodes): do not read a failed transport as a dead backend Over a tunnel, a probe that cannot get a stream to a backend fails the same way as a probe of a dead backend, with codes.Unavailable. Four places acted on that code by removing the rows of a loaded model or by shutting the model down on its worker: - the probe of the reconciler (a new outcome, ProbeUnknown, is neither a vote to reap nor a proof of life, and leaves the streak as it was); - the per-model check of the health monitor; - the warm path of the router, which now keeps the row, gives back the reservation and falls through to a load where a route exists; - the eviction of a remote model after a connection error, and the check of a cached remote model in the model loader. A refusal that the worker wrote about the backend is still evidence, and a dial with no custom dialer behaves as before. The decorators of a client implement Unwrap, so that the answer is visible through them. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/file_staging_client.go | 4 + core/services/nodes/health.go | 9 + core/services/nodes/health_mock_test.go | 5 + core/services/nodes/inflight.go | 7 + core/services/nodes/probe_cache.go | 39 +++- core/services/nodes/reconciler.go | 14 ++ core/services/nodes/router.go | 31 ++- core/services/nodes/router_test.go | 15 +- core/services/nodes/transport_failure_test.go | 203 ++++++++++++++++++ pkg/model/connection_evicting_client.go | 12 ++ ...onnection_evicting_client_internal_test.go | 51 +++++ pkg/model/loader.go | 9 +- pkg/model/loader_test.go | 35 +++ 13 files changed, 418 insertions(+), 16 deletions(-) create mode 100644 core/services/nodes/transport_failure_test.go create mode 100644 pkg/model/connection_evicting_client_internal_test.go diff --git a/core/services/nodes/file_staging_client.go b/core/services/nodes/file_staging_client.go index e03671a154da..3a1b33193485 100644 --- a/core/services/nodes/file_staging_client.go +++ b/core/services/nodes/file_staging_client.go @@ -42,6 +42,10 @@ type FileStagingClient struct { remoteModelPath string // set during LoadModel from staged ModelPath } +// Unwrap returns the client that this one decorates, so that a caller can ask it +// why a dial failed. See grpc.BackendUnwrapper. +func (c *FileStagingClient) Unwrap() grpc.Backend { return c.Backend } + type ttsReference struct { Audio string `json:"audio"` Text string `json:"text"` diff --git a/core/services/nodes/health.go b/core/services/nodes/health.go index 970dfda0ddbf..ad9f235f0cad 100644 --- a/core/services/nodes/health.go +++ b/core/services/nodes/health.go @@ -9,6 +9,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/advisorylock" + grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/xlog" "gorm.io/gorm" ) @@ -193,9 +194,17 @@ func (hm *HealthMonitor) doCheckAll(ctx context.Context) { mCheckCtx, mCancel := context.WithTimeout(ctx, 5*time.Second) ok, _ := mClient.HealthCheck(mCheckCtx) mCancel() + // A dial that failed in the transport says nothing about the + // backend. It is not a miss, and it does not clear a streak either. + transportFailed := !ok && grpc.TransportFailureOf(mClient) != nil if closer, ok := mClient.(io.Closer); ok { closer.Close() } + if transportFailed { + xlog.Debug("Model backend probe could not reach the backend, not counting it", + "node", node.ID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address) + continue + } key := modelKey{NodeID: node.ID, ModelName: m.ModelName, ReplicaIndex: m.ReplicaIndex} hm.missesMu.Lock() diff --git a/core/services/nodes/health_mock_test.go b/core/services/nodes/health_mock_test.go index fc00e6bb1b53..1196ed94c84c 100644 --- a/core/services/nodes/health_mock_test.go +++ b/core/services/nodes/health_mock_test.go @@ -133,8 +133,13 @@ func (f *fakeNodeHealthStore) RemoveNodeModel(_ context.Context, nodeID, modelNa type fakeBackendClient struct { healthy bool err error + // dialErr is what LastDialError reports: the error of the custom dialer of a + // client whose transport failed. + dialErr error } +func (c *fakeBackendClient) LastDialError() error { return c.dialErr } + func (c *fakeBackendClient) IsBusy() bool { return false } func (c *fakeBackendClient) HealthCheck(_ context.Context) (bool, error) { return c.healthy, c.err diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index 08df8411a5d2..7831a806ba21 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -56,6 +56,13 @@ func NewInFlightTrackingClient(inner grpc.Backend, registry InFlightTracker, nod } } +// Unwrap returns the client that this one decorates, so that a caller can ask it +// why a dial failed. See grpc.BackendUnwrapper. +func (c *InFlightTrackingClient) Unwrap() grpc.Backend { + backend, _ := c.ControlBackend.(grpc.Backend) + return backend +} + // OnFirstComplete registers a callback that fires once after the first tracked // inference call completes. This is used to release the initial in-flight // reservation (set during model load) after the triggering request finishes, diff --git a/core/services/nodes/probe_cache.go b/core/services/nodes/probe_cache.go index 422e36ede4e2..32f3a3d0ccf2 100644 --- a/core/services/nodes/probe_cache.go +++ b/core/services/nodes/probe_cache.go @@ -73,22 +73,49 @@ func (c *probeCache) Invalidate(key string) { // probes invalidate the cache, so a transient miss doesn't pin every // subsequent request to a re-probe. func (c *probeCache) DoOrCached(key string, probe func() bool) bool { + return c.DoOrCachedVerdict(key, func() probeVerdict { + if probe() { + return probeAlive + } + return probeDead + }) == probeAlive +} + +// probeVerdict is what a probe of a backend found. Only probeDead is a statement +// that the backend is gone. +type probeVerdict int + +const ( + // probeAlive: the backend answered. + probeAlive probeVerdict = iota + // probeDead: the backend did not answer, and the probe reached its host. + probeDead + // probeUnknown: the probe could not reach the host, because the transport + // failed. Nothing was learned about the backend. + probeUnknown +) + +// DoOrCachedVerdict is DoOrCached for a probe that can say that it learned +// nothing. Only probeAlive is cached, and every other verdict invalidates the +// cache. Callers that coalesce on the same key all get the verdict of the one +// probe that ran. +func (c *probeCache) DoOrCachedVerdict(key string, probe func() probeVerdict) probeVerdict { if c.IsFresh(key) { - return true + return probeAlive } v, _, _ := c.flight.Do(key, func() (any, error) { // Double-check after potentially waiting: another caller in this // flight may have just populated the cache. if c.IsFresh(key) { - return true, nil + return probeAlive, nil } - ok := probe() - if ok { + verdict := probe() + if verdict == probeAlive { c.markFresh(key) } else { c.Invalidate(key) } - return ok, nil + return verdict, nil }) - return v.(bool) + return v.(probeVerdict) } diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index 7b2a2d9ccb42..e958f8e63fd6 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -12,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/workerctl" + grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/xlog" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -33,6 +34,10 @@ const ( // ProbeUnreachable: nothing is listening (connection refused), or the // backend answered and affirmatively reported itself unhealthy. ProbeUnreachable + // ProbeUnknown: the probe could not get a stream to the backend, because + // the transport failed. Nothing was learned about the backend, so it is not + // a vote to reap and not a proof of life. + ProbeUnknown ) // ModelProber checks the state of a model's backend process. @@ -67,6 +72,9 @@ func (g grpcModelProber) Probe(ctx context.Context, nodeID, address string) Prob probeCtx, cancel := context.WithTimeout(ctx, probeTimeout) defer cancel() ok, err := client.HealthCheck(probeCtx) + if !ok && grpc.TransportFailureOf(client) != nil { + return ProbeUnknown + } return classifyProbeOutcome(ok, err) } @@ -493,6 +501,12 @@ func (rc *ReplicaReconciler) probeLoadedModels(ctx context.Context) { _ = rc.registry.db.WithContext(ctx).Model(&NodeModel{}). Where("id = ?", m.ID).Update("updated_at", time.Now()).Error continue + case ProbeUnknown: + // The transport failed. This is neither a vote to reap nor proof of + // life, so the streak stays as it is. + xlog.Debug("Reconciler: model probe could not reach the backend, not counting it", + "node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address) + continue case ProbeBusy: // Reachable but mid-request. Proof of life, so clear the streak. rc.clearProbeFailures(m.ID) diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 2ea18bc195a2..e28b679cdce6 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -776,7 +776,20 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route replicaIdx := nm.ReplicaIndex // Verify the backend process is still alive via gRPC health check - if !r.probeHealth(ctx, node, modelAddr) { + switch r.probeHealth(ctx, node, modelAddr) { + case probeUnknown: + // The transport failed, so nothing is known about the backend. The row + // stays: a replica that is loaded on a worker that cannot be reached from + // here for the moment is not a replica that is gone. This request goes + // to a cold load, which places the model where a route exists. + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release routing reservation of an unreachable replica", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + xlog.Warn("Backend not reachable through the transport, keeping its row and falling through to a load elsewhere", + "node", node.Name, "model", att.modelName, "replica", replicaIdx) + return nil + case probeDead: // Stale — roll back the increment, remove the specific replica row, fall through if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { xlog.Warn("Failed to release stale routing reservation", @@ -2043,15 +2056,25 @@ func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir // burst of N requests for a cold cache costs at most one round-trip, not N. // Failed probes invalidate the cache so the staleness recovery path // (DecrementInFlight + RemoveNodeModel) still triggers on the next request. -func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr string) bool { +// +// The verdict is three-valued. A probe that could not get a stream to the backend +// because the transport failed learned nothing, and it must not read as a dead +// backend: the caller would remove the row of a model that is loaded. +func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr string) probeVerdict { key := node.ID + "|" + addr - return r.probeCache.DoOrCached(key, func() bool { + return r.probeCache.DoOrCachedVerdict(key, func() probeVerdict { client := r.buildClientForAddr(node, addr, false) defer closeClient(client) checkCtx, cancel := context.WithTimeout(ctx, 2*time.Second) defer cancel() ok, _ := client.HealthCheck(checkCtx) - return ok + switch { + case ok: + return probeAlive + case grpc.TransportFailureOf(client) != nil: + return probeUnknown + } + return probeDead }) } diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index e5a972050f5b..20f586929bc0 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -499,17 +499,22 @@ type stubBackend struct { healthResult bool healthErr error - loadResult *pb.Result - loadErr error - loadHook func(*pb.ModelOptions) - loadOpts []*pb.ModelOptions - mu sync.Mutex + // dialErr is what LastDialError reports: the error of the custom dialer of + // a client whose transport failed. + dialErr error + loadResult *pb.Result + loadErr error + loadHook func(*pb.ModelOptions) + loadOpts []*pb.ModelOptions + mu sync.Mutex } func (f *stubBackend) HealthCheck(_ context.Context) (bool, error) { return f.healthResult, f.healthErr } +func (f *stubBackend) LastDialError() error { return f.dialErr } + func (f *stubBackend) LoadModel(_ context.Context, opts *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { cloned := proto.Clone(opts).(*pb.ModelOptions) f.mu.Lock() diff --git a/core/services/nodes/transport_failure_test.go b/core/services/nodes/transport_failure_test.go new file mode 100644 index 000000000000..85146703725b --- /dev/null +++ b/core/services/nodes/transport_failure_test.go @@ -0,0 +1,203 @@ +package nodes + +import ( + "context" + "errors" + "fmt" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +// A dial that fails in the transport says nothing about the backend. The tunnel +// carrier makes that the ordinary way for a probe to fail: a worker that is not +// connected to any replica right now cannot be probed, and it is well. Every +// place that reaps a row on a failed probe must tell the two apart. + +// backendAnswer is the error of a dialer that reached the host of a backend and +// was told that the backend is gone. +type backendAnswer struct{ error } + +func (backendAnswer) IsBackendAnswer() bool { return true } + +var errNoTunnel = errors.New("no tunnel is held for that node") + +var _ = Describe("A probe that could not reach the backend", func() { + unavailable := status.Error(codes.Unavailable, "dial failed") + + DescribeTable("is classified by what failed", + func(mk func() *fakeBackendClient, want ProbeOutcome) { + f := &recordingFactory{next: func() grpc.Backend { return mk() }} + Expect(grpcModelProber{clients: f}.Probe(context.Background(), "n1", "a:1")).To(Equal(want)) + }, + Entry("the transport failed: nothing was learned", func() *fakeBackendClient { + return &fakeBackendClient{err: unavailable, dialErr: errNoTunnel} + }, ProbeUnknown), + Entry("the transport failed and the call timed out: still nothing learned", func() *fakeBackendClient { + return &fakeBackendClient{err: context.DeadlineExceeded, dialErr: errNoTunnel} + }, ProbeUnknown), + Entry("the host answered that the backend is gone: evidence", func() *fakeBackendClient { + return &fakeBackendClient{err: unavailable, dialErr: backendAnswer{errors.New("the worker could not reach the backend")}} + }, ProbeUnreachable), + Entry("the backend is well, whatever the last dial said", func() *fakeBackendClient { + return &fakeBackendClient{healthy: true, dialErr: errNoTunnel} + }, ProbeAlive), + Entry("no custom dialer: a refused connection is still death", func() *fakeBackendClient { + return &fakeBackendClient{err: unavailable} + }, ProbeUnreachable), + ) +}) + +var _ = Describe("The reconciler and a probe that learned nothing", func() { + It("never reaps a replica on it, however often it repeats", func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db := testutil.SetupTestDB() + registry, err := NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + node := &BackendNode{Name: "tunnel-node", NodeType: NodeTypeBackend} + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + Expect(db.Create(&NodeModel{ + ID: "m-1", NodeID: node.ID, ModelName: "m-1", Address: "127.0.0.1:50052", + State: "loaded", UpdatedAt: time.Now().Add(-5 * time.Minute), + }).Error).To(Succeed()) + + prober := &fakeProber{outcomes: map[string]ProbeOutcome{"127.0.0.1:50052": ProbeUnknown}} + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db, Prober: prober, ProbeStaleAfter: 2 * time.Minute}) + for range probeFailuresBeforeReap * 3 { + rc.probeLoadedModels(context.Background()) + Expect(db.Model(&NodeModel{}).Where("id = ?", "m-1").Update("updated_at", time.Now().Add(-5*time.Minute)).Error).To(Succeed()) + } + var after NodeModel + Expect(db.First(&after, "id = ?", "m-1").Error).To(Succeed(), "a transport that failed is not a dead backend") + }) + + It("keeps a streak of real failures where it was, and does not count the unknown one", func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db := testutil.SetupTestDB() + registry, err := NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + node := &BackendNode{Name: "n", NodeType: NodeTypeBackend} + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + Expect(db.Create(&NodeModel{ + ID: "m-2", NodeID: node.ID, ModelName: "m-2", Address: "127.0.0.1:50053", + State: "loaded", UpdatedAt: time.Now().Add(-5 * time.Minute), + }).Error).To(Succeed()) + prober := &fakeProber{outcomes: map[string]ProbeOutcome{"127.0.0.1:50053": ProbeUnreachable}} + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db, Prober: prober, ProbeStaleAfter: 2 * time.Minute}) + stale := func() { + Expect(db.Model(&NodeModel{}).Where("id = ?", "m-2").Update("updated_at", time.Now().Add(-5*time.Minute)).Error).To(Succeed()) + } + + // One miss short of the threshold, then a pass that learns nothing, then + // the last miss: the unknown pass neither wipes nor adds to the streak. + for range probeFailuresBeforeReap - 1 { + rc.probeLoadedModels(context.Background()) + stale() + } + prober.outcomes["127.0.0.1:50053"] = ProbeUnknown + rc.probeLoadedModels(context.Background()) + stale() + var mid NodeModel + Expect(db.First(&mid, "id = ?", "m-2").Error).To(Succeed()) + + prober.outcomes["127.0.0.1:50053"] = ProbeUnreachable + rc.probeLoadedModels(context.Background()) + var after NodeModel + Expect(db.First(&after, "id = ?", "m-2").Error).ToNot(Succeed(), "the streak was intact, so this miss reaps") + }) +}) + +var _ = Describe("The health monitor and a probe that learned nothing", func() { + build := func() (*fakeNodeHealthStore, *fakeBackendClientFactory, *HealthMonitor) { + store := newFakeNodeHealthStore() + factory := newFakeBackendClientFactory() + hm := newTestHealthMonitor(store, factory, true, 5*time.Minute) + hm.perModelHealthCheck = true + node := makeTestNode("n", "tunnel-worker", "127.0.0.1:50051", StatusHealthy, time.Now()) + store.addNode(node) + store.addNodeModel("n", NodeModel{NodeID: "n", ModelName: "m", Address: "127.0.0.1:50053"}) + return store, factory, hm + } + + It("does not count it as a miss, however often it repeats", func() { + store, factory, hm := build() + factory.setClient("127.0.0.1:50053", &fakeBackendClient{err: fmt.Errorf("unavailable"), dialErr: errNoTunnel}) + for range perModelMissThreshold * 3 { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) + }) + + It("still reaps after the threshold when the host answered that the backend is gone", func() { + store, factory, hm := build() + factory.setClient("127.0.0.1:50053", &fakeBackendClient{err: fmt.Errorf("unavailable"), dialErr: backendAnswer{errors.New("gone")}}) + for range perModelMissThreshold { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).To(ContainElement("RemoveNodeModel:n:m:0")) + }) +}) + +var _ = Describe("The decorators of a client", func() { + It("let a caller see the dial error of the client under them", func() { + raw := &fakeBackendClient{dialErr: errNoTunnel} + staged := NewFileStagingClient(raw, nil, "n") + tracked := NewInFlightTrackingClient(staged, nil, "n", "m", 0) + + Expect(errors.Is(grpc.LastDialErrorOf(staged), errNoTunnel)).To(BeTrue()) + Expect(errors.Is(grpc.LastDialErrorOf(tracked), errNoTunnel)).To(BeTrue()) + Expect(errors.Is(grpc.TransportFailureOf(tracked), errNoTunnel)).To(BeTrue()) + }) +}) + +var _ = Describe("The router and a replica it cannot reach", func() { + cached := func() (*fakeModelRouter, *BackendNode) { + old := &BackendNode{ID: "n-old", Name: "old-node", Address: "10.0.0.70:50051"} + fresh := &BackendNode{ID: "n-new", Name: "new-node", Address: "10.0.0.71:50051"} + return &fakeModelRouter{ + findAndLockNode: old, + findAndLockNM: &NodeModel{NodeID: "n-old", ModelName: "m", Address: "10.0.0.70:9001"}, + findIdleNode: fresh, + }, fresh + } + route := func(reg *fakeModelRouter, backend *stubBackend) { + GinkgoHelper() + backend.loadResult = &pb.Result{Success: true} + unloader := &fakeUnloader{installReply: &workerctl.BackendInstallReply{Success: true, Address: "10.0.0.71:9001"}} + router := NewSmartRouter(reg, SmartRouterOptions{Unloader: unloader, ClientFactory: &stubClientFactory{client: backend}}) + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) + Expect(err).ToNot(HaveOccurred()) + } + + It("keeps the row of a replica whose probe could not get a stream, and gives its reservation back", func() { + reg, _ := cached() + route(reg, &stubBackend{healthResult: false, healthErr: errors.New("unavailable"), dialErr: errNoTunnel}) + Expect(reg.removeCalls).ToNot(ContainElement("n-old:m"), "a worker that cannot be reached from here is not a replica that is gone") + Expect(reg.decrementCalls).To(ContainElement("n-old:m")) + }) + + It("removes the row of a replica whose backend does not answer on a reachable host", func() { + reg, _ := cached() + route(reg, &stubBackend{healthResult: false, healthErr: errors.New("unavailable")}) + Expect(reg.removeCalls).To(ContainElement("n-old:m")) + }) + + It("removes the row when the host answered that the backend is gone", func() { + reg, _ := cached() + route(reg, &stubBackend{healthResult: false, healthErr: errors.New("unavailable"), dialErr: backendAnswer{errors.New("gone")}}) + Expect(reg.removeCalls).To(ContainElement("n-old:m")) + }) +}) diff --git a/pkg/model/connection_evicting_client.go b/pkg/model/connection_evicting_client.go index 5cf920ede6d2..39fd48d98732 100644 --- a/pkg/model/connection_evicting_client.go +++ b/pkg/model/connection_evicting_client.go @@ -30,8 +30,20 @@ func newConnectionEvictingClient(inner grpc.Backend, modelID string, evict func( } } +// Unwrap returns the client that this one decorates, so that a caller can ask it +// why a dial failed. See grpc.BackendUnwrapper. +func (c *ConnectionEvictingClient) Unwrap() grpc.Backend { return c.Backend } + func (c *ConnectionEvictingClient) checkErr(err error) { if err != nil && isConnectionError(err) { + // A dial that failed in the transport says nothing about the backend. The + // eviction shuts the model down on the worker and removes its rows, and a + // backend that could not be reached has not been shown to be gone. + if transportErr := grpc.TransportFailureOf(c.Backend); transportErr != nil { + xlog.Debug("Connection error during inference came from the transport; keeping the model", + "model", c.modelID, "error", transportErr) + return + } c.once.Do(func() { xlog.Warn("Connection error during inference, evicting model from cache", "model", c.modelID, "error", err) diff --git a/pkg/model/connection_evicting_client_internal_test.go b/pkg/model/connection_evicting_client_internal_test.go new file mode 100644 index 000000000000..18ace82afb2f --- /dev/null +++ b/pkg/model/connection_evicting_client_internal_test.go @@ -0,0 +1,51 @@ +package model + +import ( + "context" + "errors" + "net" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +type workerAnswer struct{ error } + +func (workerAnswer) IsBackendAnswer() bool { return true } + +// A connection error during an inference shuts the model down on the worker and +// removes its rows. A dial that failed in the transport is a connection error to +// gRPC and says nothing about the backend. +var _ = Describe("ConnectionEvictingClient", func() { + evictedBy := func(dialErr error) bool { + GinkgoHelper() + var evicted atomic.Bool + inner := grpc.NewClientWithDialer("10.0.0.5:50051", false, nil, false, "", func(context.Context, string) (net.Conn, error) { + return nil, dialErr + }) + client := newConnectionEvictingClient(inner, "m", func() { evicted.Store(true) }) + _, err := client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(err).To(HaveOccurred()) + return evicted.Load() + } + + It("keeps the model when the transport could not get a stream to the backend", func() { + Expect(evictedBy(errors.New("no tunnel is held for that node"))).To(BeFalse()) + }) + + It("evicts the model when the worker answered that the backend is gone", func() { + Expect(evictedBy(workerAnswer{errors.New("the worker could not reach the backend")})).To(BeTrue()) + }) + + It("lets a caller see through it to the dial error of the client", func() { + boom := errors.New("no route") + inner := grpc.NewClientWithDialer("x:1", false, nil, false, "", func(context.Context, string) (net.Conn, error) { return nil, boom }) + _, _ = inner.Predict(context.Background(), &pb.PredictOptions{}) + client := newConnectionEvictingClient(inner, "m", func() {}) + Expect(errors.Is(grpc.LastDialErrorOf(client), boom)).To(BeTrue()) + }) +}) diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 85e39f3b57b8..616e619733e2 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -12,6 +12,7 @@ import ( "sync/atomic" "time" + grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/utils" @@ -716,7 +717,13 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { if process == nil { // Remote/distributed model — no local process to check. // Only evict on definitive connection errors (node is down). - // Timeouts may mean the node is busy, so keep the model cached. + // Timeouts may mean the node is busy, so keep the model cached. So + // does a dial that failed in the transport: it says nothing about + // the backend. + if transportErr := grpc.TransportFailureOf(client); transportErr != nil { + xlog.Warn("Remote model health check failed in the transport, keeping cached", "model", s, "error", transportErr) + return m + } if isConnectionError(err) { xlog.Warn("Remote model unreachable (connection error), removing from cache", "model", s, "error", err) if delErr := ml.deleteProcess(cTimeout, s, false); delErr != nil { diff --git a/pkg/model/loader_test.go b/pkg/model/loader_test.go index 511169e38152..469e0af3b9eb 100644 --- a/pkg/model/loader_test.go +++ b/pkg/model/loader_test.go @@ -3,6 +3,7 @@ package model_test import ( "context" "errors" + "net" "os" "path/filepath" "sync" @@ -201,6 +202,34 @@ var _ = Describe("ModelLoader", func() { }) }) + Context("Remote model whose transport failed", func() { + // A client with its own dialer, as a tunnel gives. Its dial fails in the + // transport and says nothing about the backend. + failing := func(dialErr error) grpcPkg.Backend { + return grpcPkg.NewClientWithDialer("10.0.0.5:50051", false, nil, false, "", func(context.Context, string) (net.Conn, error) { + return nil, dialErr + }) + } + + load := func(id string, client grpcPkg.Backend) { + GinkgoHelper() + remote := model.NewModelWithClient(id, "10.0.0.5:50051", client) + _, err := modelLoader.LoadModel(id, "test.model", func(string, string, string) (*model.Model, error) { return remote, nil }) + Expect(err).ToNot(HaveOccurred()) + } + + It("keeps the model cached when the transport could not reach the worker", func() { + load("unreachable", failing(errors.New("no tunnel is held for that node"))) + Expect(modelLoader.CheckIsLoaded("unreachable")).ToNot(BeNil(), + "a worker that cannot be reached from here has not been shown to have lost the model") + }) + + It("evicts the model when the worker answered that the backend is gone", func() { + load("gone", failing(backendGone{errors.New("the worker could not reach the backend")})) + Expect(modelLoader.CheckIsLoaded("gone")).To(BeNil()) + }) + }) + Context("ShutdownModel", func() { It("should shutdown a loaded model", func() { mockLoader := func(modelID, modelName, modelFile string) (*model.Model, error) { @@ -616,3 +645,9 @@ var _ = Describe("ModelLoader", func() { }) }) }) + +// backendGone is the error of a dialer whose host answered that the backend is +// gone: evidence about the backend, like a refused connection. +type backendGone struct{ error } + +func (backendGone) IsBackendAnswer() bool { return true } From 12d40e28c437bf230aac71332caee774c5a39e59 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:12:50 +0000 Subject: [PATCH 063/162] feat(nodes): reach the backends and the HTTP server of a worker through a dialer The pieces that build their own connection to a worker now go through the dial seam, so that a worker that holds a tunnel and has no address can be reached by all of them. - NewDialerClientFactory builds gRPC clients whose backend is reached by the dialer of the node. - The HTTP file stager reports a missing dialer as ErrNoRoute and not as a panic, and ForgetNode drops the client of a node that left and closes its idle streams. - A worker with no address gets a reserved host name for the URL of its HTTP server. The name is never dialled. The proxy of the environment is not used for it, because the transport would dial the proxy through the tunnel. - The backend-logs proxy uses that host, and answers an older worker that has a gRPC address and no HTTP server as before. - The models that DistributedModelStore.Range builds from the registry carry a client from the factory, where they carried none and a caller got a TCP dial to the address of the node. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 4 + core/application/startup.go | 1 + core/http/endpoints/localai/nodes.go | 33 ++++- .../localai/nodes_logs_tunnel_test.go | 92 +++++++++++++ .../nodes/client_factory_dialer_test.go | 127 ++++++++++++++++++ core/services/nodes/control_client.go | 13 +- core/services/nodes/distributed_store.go | 33 ++++- core/services/nodes/distributed_store_test.go | 30 +++++ core/services/nodes/file_stager_http.go | 73 ++++++++-- core/services/nodes/interfaces.go | 26 ++++ 10 files changed, 410 insertions(+), 22 deletions(-) create mode 100644 core/http/endpoints/localai/nodes_logs_tunnel_test.go create mode 100644 core/services/nodes/client_factory_dialer_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 22309eaf724a..6bb7d226d3d1 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -61,6 +61,9 @@ type DistributedServices struct { ModelAdapter *nodes.ModelRouterAdapter Unloader nodes.NodeControl ModelCleanup *nodes.ModelCleanupService + // Clients builds the gRPC client for a backend of a worker through the + // carrier that is active. + Clients nodes.BackendClientFactory // WorkerHTTPDial reaches a worker's own HTTP server for the admin // backend-logs proxy, the same way the HTTP file stager does. @@ -621,6 +624,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade ModelAdapter: modelAdapter, Unloader: remoteUnloader, ModelCleanup: modelCleanup, + Clients: clientFactory, WorkerHTTPDial: workerHTTPDial, diff --git a/core/application/startup.go b/core/application/startup.go index e77d4e5ba638..c3e8ffcdc83b 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -306,6 +306,7 @@ func New(opts ...config.AppOption) (*Application, error) { distStore := nodes.NewDistributedModelStore( model.NewInMemoryModelStore(), distSvc.Registry, + nodes.WithClientFactory(distSvc.Clients), ) application.modelLoader.SetModelStore(distStore) // Drop the local stub when a model's last replica leaves the registry. diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index ff5d19891943..acee0624bb00 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -977,11 +977,12 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { + host, ok := workerLogsHost(node) + if !ok { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) } - resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, node.HTTPAddress, "/v1/backend-logs", registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, host, "/v1/backend-logs", registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -1007,12 +1008,13 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { + host, ok := workerLogsHost(node) + if !ok { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) } path := "/v1/backend-logs/" + url.PathEscape(modelID) - resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, node.HTTPAddress, path, registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, host, path, registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -1060,7 +1062,7 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s } // Dial the worker WebSocket - workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", node.HTTPAddress, url.PathEscape(modelID)) + workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", nodes.WorkerHTTPHost(node.ID, node.HTTPAddress), url.PathEscape(modelID)) workerHeaders := http.Header{} if registrationToken != "" { workerHeaders.Set("Authorization", "Bearer "+registrationToken) @@ -1536,6 +1538,21 @@ func DeleteSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { } } +// workerLogsHost returns the host of the URL of a request to the HTTP server of a +// worker, and false when the node has no HTTP server to ask. +// +// A worker that holds a tunnel registers no address at all: nothing dials it, and +// the dialer opens a stream on its tunnel whatever host the URL names. Its host +// is the reserved name of nodes.WorkerHTTPHost. A worker that registered a +// gRPC address and no HTTP address is an older worker with no log endpoint, and +// the answer for it has not changed. +func workerLogsHost(node *nodes.BackendNode) (string, bool) { + if node.HTTPAddress == "" && node.Address != "" { + return "", false + } + return nodes.WorkerHTTPHost(node.ID, node.HTTPAddress), true +} + // proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer token auth. // The connection goes through dialFor(nodeID) because the advertised address // alone does not say how this frontend reaches that worker. @@ -1554,6 +1571,12 @@ func proxyHTTPToWorker(ctx context.Context, dialFor nodes.WorkerNetDialerFor, no t := httpclient.HardenedTransport() t.DialContext = dialFor(nodeID) + // A worker without an address is reached through a tunnel, and a proxy of the + // environment cannot carry that: the transport would dial the proxy through + // the dialer of the worker. The proxy stays for every other worker. + if nodes.IsTunnelOnlyHost(httpAddress) { + t.Proxy = nil + } client := httpclient.NewWithTimeout(15*time.Second, httpclient.WithTransport(t)) return client.Do(req) } diff --git a/core/http/endpoints/localai/nodes_logs_tunnel_test.go b/core/http/endpoints/localai/nodes_logs_tunnel_test.go new file mode 100644 index 000000000000..09187f28b82e --- /dev/null +++ b/core/http/endpoints/localai/nodes_logs_tunnel_test.go @@ -0,0 +1,92 @@ +package localai + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/nodes" +) + +var _ = Describe("Backend logs of a worker that holds a tunnel", func() { + Describe("the host of the request", func() { + It("names the reserved host for a worker with no address at all", func() { + host, ok := workerLogsHost(&nodes.BackendNode{ID: "n1"}) + Expect(ok).To(BeTrue()) + Expect(host).To(Equal("n1.worker.invalid")) + Expect(nodes.IsTunnelOnlyHost(host)).To(BeTrue()) + }) + + It("keeps the HTTP address of a worker that has one", func() { + host, ok := workerLogsHost(&nodes.BackendNode{ID: "n1", Address: "10.0.0.1:50051", HTTPAddress: "10.0.0.1:50050"}) + Expect(ok).To(BeTrue()) + Expect(host).To(Equal("10.0.0.1:50050")) + }) + + It("still refuses a worker that has a gRPC address and no HTTP server, as before", func() { + _, ok := workerLogsHost(&nodes.BackendNode{ID: "n1", Address: "10.0.0.1:50051"}) + Expect(ok).To(BeFalse()) + }) + }) + + Describe("the proxy of the environment", func() { + var ( + srv *httptest.Server + mu sync.Mutex + dials []string + ) + + BeforeEach(func() { + dials = nil + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.WriteString(w, "logs for "+r.Host) + })) + DeferCleanup(srv.Close) + // A proxy that nothing listens on. A request that honours it fails. + GinkgoT().Setenv("HTTP_PROXY", "http://127.0.0.1:1") + GinkgoT().Setenv("http_proxy", "http://127.0.0.1:1") + GinkgoT().Setenv("NO_PROXY", "") + GinkgoT().Setenv("no_proxy", "") + }) + + dialer := func() nodes.WorkerNetDialerFor { + return func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, addr string) (net.Conn, error) { + mu.Lock() + dials = append(dials, addr) + mu.Unlock() + var d net.Dialer + return d.DialContext(ctx, "tcp", srv.Listener.Addr().String()) + } + } + } + + It("is not used to reach a worker through its tunnel, so the dialer gets the stream", func() { + resp, err := proxyHTTPToWorker(context.Background(), dialer(), "n1", "n1.worker.invalid", "/v1/backend-logs", "token") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + body, _ := io.ReadAll(resp.Body) + Expect(string(body)).To(Equal("logs for n1.worker.invalid")) + mu.Lock() + defer mu.Unlock() + Expect(dials).To(HaveLen(1)) + Expect(dials[0]).To(HavePrefix("n1.worker.invalid"), "the dialer was asked for the worker, not for a proxy") + }) + + It("is still used for a worker that has an address, as before", func() { + _, err := proxyHTTPToWorker(context.Background(), dialer(), "n1", "10.0.0.1:50050", "/v1/backend-logs", "token") + Expect(err).ToNot(HaveOccurred()) + mu.Lock() + defer mu.Unlock() + Expect(dials).To(HaveLen(1)) + Expect(strings.HasPrefix(dials[0], "127.0.0.1:1")).To(BeTrue(), "with a proxy in the environment the dial goes to the proxy: %v", dials) + }) + }) +}) diff --git a/core/services/nodes/client_factory_dialer_test.go b/core/services/nodes/client_factory_dialer_test.go new file mode 100644 index 000000000000..330b996960e9 --- /dev/null +++ b/core/services/nodes/client_factory_dialer_test.go @@ -0,0 +1,127 @@ +package nodes + +import ( + "context" + "errors" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + ggrpc "google.golang.org/grpc" +) + +var _ = Describe("The client factory with a dialer", func() { + It("hands the dialer of the node to the client, and the node id selects it", func() { + var asked []string + boom := errors.New("no route to that node") + factory := NewDialerClientFactory("", func(nodeID string) func(context.Context, string) (net.Conn, error) { + asked = append(asked, nodeID) + return func(context.Context, string) (net.Conn, error) { return nil, boom } + }) + + client := factory.NewClient("node-7", "127.0.0.1:50051", false) + ok, err := client.HealthCheck(context.Background()) + Expect(ok).To(BeFalse()) + Expect(err).To(HaveOccurred()) + Expect(asked).To(Equal([]string{"node-7"})) + Expect(errors.Is(grpc.TransportFailureOf(client), boom)).To(BeTrue()) + }) + + It("reaches a backend that listens on a socket the dialer chose", func() { + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = lis.Close() }) + stop := serveBackend(lis) + DeferCleanup(stop) + + var dials atomic.Int32 + factory := NewDialerClientFactory("", func(string) func(context.Context, string) (net.Conn, error) { + return func(ctx context.Context, addr string) (net.Conn, error) { + dials.Add(1) + Expect(addr).To(Equal("a-name-for-a-process:1")) + var d net.Dialer + return d.DialContext(ctx, "tcp", lis.Addr().String()) + } + }) + client := factory.NewClient("n", "a-name-for-a-process:1", false) + ok, err := client.HealthCheck(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(dials.Load()).To(BeNumerically(">=", 1)) + }) +}) + +type healthServer struct{ pb.UnimplementedBackendServer } + +func (healthServer) Health(context.Context, *pb.HealthMessage) (*pb.Reply, error) { + return &pb.Reply{Message: []byte("OK")}, nil +} + +// serveBackend starts a backend that answers its health check on lis. +func serveBackend(lis net.Listener) func() { + srv := ggrpc.NewServer() + pb.RegisterBackendServer(srv, healthServer{}) + go func() { _ = srv.Serve(lis) }() + return srv.Stop +} + +var _ = Describe("The HTTP file stager without a route", func() { + It("reports a missing dialer as ErrNoRoute and not as a panic", func() { + stager := NewHTTPFileStager(func(string) (string, error) { return "w.worker.invalid", nil }, "", nil) + _, err := stager.ListRemoteDir(context.Background(), "n", "k") + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue(), "%v", err) + }) + + It("reports a node that has no dialer the same way", func() { + stager := NewHTTPFileStager(func(string) (string, error) { return "w.worker.invalid", nil }, "", func(string) func(context.Context, string, string) (net.Conn, error) { return nil }) + _, err := stager.ListRemoteDir(context.Background(), "n", "k") + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue(), "%v", err) + }) + + It("reports a dialer that finds no route as that error, with the chain intact", func() { + stager := NewHTTPFileStager(func(string) (string, error) { return "w.worker.invalid", nil }, "", func(string) func(context.Context, string, string) (net.Conn, error) { + return func(context.Context, string, string) (net.Conn, error) { return nil, ErrNoRoute } + }) + _, err := stager.ListRemoteDir(context.Background(), "n", "k") + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue()) + }) +}) + +var _ = Describe("The HTTP file stager and a node that left", func() { + It("closes the idle streams of the node and dials again after ForgetNode", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"files":[]}`)) + })) + DeferCleanup(srv.Close) + var dials atomic.Int32 + stager := NewHTTPFileStager(func(string) (string, error) { return strings.TrimPrefix(srv.URL, "http://"), nil }, "", + func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + dials.Add(1) + var d net.Dialer + return d.DialContext(ctx, "tcp", srv.Listener.Addr().String()) + } + }) + for range 3 { + _, err := stager.ListRemoteDir(context.Background(), "n", "k") + Expect(err).ToNot(HaveOccurred()) + } + Expect(dials.Load()).To(Equal(int32(1))) + + stager.ForgetNode("n") + _, err := stager.ListRemoteDir(context.Background(), "n", "k") + Expect(err).ToNot(HaveOccurred()) + Eventually(dials.Load, time.Second).Should(Equal(int32(2))) + + var none *HTTPFileStager + none.ForgetNode("n") + }) +}) diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go index cdf83135f948..850563fc02f5 100644 --- a/core/services/nodes/control_client.go +++ b/core/services/nodes/control_client.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "net/http" + "strings" "sync" "time" @@ -62,7 +63,17 @@ func WorkerHTTPHost(nodeID, addr string) string { if addr != "" { return addr } - return nodeID + ".worker.invalid" + return nodeID + tunnelOnlyHostSuffix +} + +// tunnelOnlyHostSuffix ends the host that WorkerHTTPHost gives to a worker that +// has no address. .invalid is reserved and resolves nowhere. +const tunnelOnlyHostSuffix = ".worker.invalid" + +// IsTunnelOnlyHost reports whether a host is one that WorkerHTTPHost made for a +// worker with no address, which only a tunnel can reach. +func IsTunnelOnlyHost(host string) bool { + return strings.HasSuffix(host, tunnelOnlyHostSuffix) } // clientFor returns the HTTP client for a node, and builds it on first use. diff --git a/core/services/nodes/distributed_store.go b/core/services/nodes/distributed_store.go index ba1379367413..c23fa05b3eb4 100644 --- a/core/services/nodes/distributed_store.go +++ b/core/services/nodes/distributed_store.go @@ -14,10 +14,37 @@ import ( type DistributedModelStore struct { local model.ModelStore registry ModelLookup + clients BackendClientFactory } -func NewDistributedModelStore(local model.ModelStore, registry ModelLookup) *DistributedModelStore { - return &DistributedModelStore{local: local, registry: registry} +// DistributedModelStoreOption changes how a DistributedModelStore is built. +type DistributedModelStoreOption func(*DistributedModelStore) + +// WithClientFactory makes the models that Range builds from the registry carry a +// client from the factory. Without it they carry none, and the first caller that +// asks one for its client gets a plain TCP dial to the address of the node. That +// is a dial that bypasses the carrier: it fails for a worker that holds a tunnel +// and has no address, and it is the one place left that builds a client outside +// the dial seam. +func WithClientFactory(f BackendClientFactory) DistributedModelStoreOption { + return func(s *DistributedModelStore) { s.clients = f } +} + +func NewDistributedModelStore(local model.ModelStore, registry ModelLookup, opts ...DistributedModelStoreOption) *DistributedModelStore { + s := &DistributedModelStore{local: local, registry: registry} + for _, opt := range opts { + opt(s) + } + return s +} + +// stub builds the model that stands for a replica that the registry holds and +// this process has not loaded. +func (s *DistributedModelStore) stub(nodeID, modelName, address string) *model.Model { + if s.clients == nil { + return model.NewModel(modelName, address, nil) + } + return model.NewModelWithClient(modelName, address, s.clients.NewClient(nodeID, address, false)) } // Get checks the local cache only. In distributed mode, models must be routed @@ -80,7 +107,7 @@ func (s *DistributedModelStore) Range(fn func(string, *model.Model) bool) { continue } - m := model.NewModel(nm.ModelName, node.Address, nil) + m := s.stub(nm.NodeID, nm.ModelName, node.Address) if !fn(nm.ModelName, m) { return } diff --git a/core/services/nodes/distributed_store_test.go b/core/services/nodes/distributed_store_test.go index 9b6e4ccc9464..cf71e55590ac 100644 --- a/core/services/nodes/distributed_store_test.go +++ b/core/services/nodes/distributed_store_test.go @@ -7,6 +7,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/LocalAI/pkg/model" ) @@ -131,3 +132,32 @@ var _ = Describe("DistributedModelStore", func() { }) }) }) + +var _ = Describe("DistributedModelStore.Range and the dial seam", func() { + It("builds the model of a replica it does not hold through the client factory", func() { + lookup := newFakeModelLookup() + lookup.nodes["node-t"] = &BackendNode{ID: "node-t", Address: ""} + lookup.allModels = []NodeModel{{NodeID: "node-t", ModelName: "tunnelled"}} + backend := &fakeBackendClient{healthy: true} + factory := &recordingFactory{next: func() grpc.Backend { return backend }} + store := NewDistributedModelStore(model.NewInMemoryModelStore(), lookup, WithClientFactory(factory)) + + var got *model.Model + store.Range(func(_ string, m *model.Model) bool { got = m; return true }) + + Expect(got).ToNot(BeNil()) + Expect(factory.calls()).To(Equal([]string{"node-t@"}), "the node id travels, because a tunnel cannot recover it from the address") + Expect(got.GRPC(false, nil)).To(BeIdenticalTo(backend), "the client comes from the seam, not from a TCP dial to the address") + }) + + It("keeps the tokenless stub of old when no factory is given", func() { + lookup := newFakeModelLookup() + lookup.nodes["node-1"] = &BackendNode{ID: "node-1", Address: "10.0.0.9:50051"} + lookup.allModels = []NodeModel{{NodeID: "node-1", ModelName: "plain"}} + store := NewDistributedModelStore(model.NewInMemoryModelStore(), lookup) + var got *model.Model + store.Range(func(_ string, m *model.Model) bool { got = m; return true }) + Expect(got).ToNot(BeNil()) + Expect(got.Process()).To(BeNil()) + }) +}) diff --git a/core/services/nodes/file_stager_http.go b/core/services/nodes/file_stager_http.go index 4cfde4bfddfc..23c681ef5849 100644 --- a/core/services/nodes/file_stager_http.go +++ b/core/services/nodes/file_stager_http.go @@ -34,7 +34,7 @@ type HTTPFileStager struct { token string dialFor WorkerNetDialerFor clientsMu sync.Mutex - clients map[string]*http.Client + clients map[string]*nodeFileClient responseTimeout time.Duration // timeout waiting for server response after upload maxRetries int // number of retry attempts for transient failures } @@ -62,24 +62,42 @@ func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token st httpAddrFor: httpAddrFor, token: token, dialFor: dialFor, - clients: map[string]*http.Client{}, + clients: map[string]*nodeFileClient{}, responseTimeout: responseTimeout, maxRetries: maxRetries, } } +// nodeFileClient is the client of one node and the transport it was built on. The +// transport is kept so that ForgetNode can close the idle streams that it holds. +type nodeFileClient struct { + client *http.Client + transport *http.Transport +} + // clientFor returns the HTTP client that reaches nodeID. Clients are per node // rather than shared because the idle pool is keyed by host:port only: two // workers reporting the same address (NAT, loopback) would otherwise be handed // each other's connections once the dialer routes by node. -func (h *HTTPFileStager) clientFor(nodeID string) *http.Client { +// +// A stager with no dialer for the node cannot reach it. That is a route that does +// not exist, and ErrNoRoute says so, where a nil function would panic at the +// first transfer. +func (h *HTTPFileStager) clientFor(nodeID string) (*http.Client, error) { h.clientsMu.Lock() defer h.clientsMu.Unlock() if c, ok := h.clients[nodeID]; ok { - return c + return c.client, nil + } + if h.dialFor == nil { + return nil, fmt.Errorf("file transfer to node %s: no dialer is configured: %w", nodeID, ErrNoRoute) + } + dial := h.dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("file transfer to node %s: no dialer for the node: %w", nodeID, ErrNoRoute) } transport := &http.Transport{ - DialContext: h.dialFor(nodeID), + DialContext: dial, ForceAttemptHTTP2: false, // HTTP/2 flow control can stall large uploads MaxIdleConns: 10, IdleConnTimeout: 90 * time.Second, @@ -94,8 +112,34 @@ func (h *HTTPFileStager) clientFor(nodeID string) *http.Client { // on the server. Instead we use ResponseHeaderTimeout on the transport // to cover only the wait-for-server-response phase. c := httpclient.New(httpclient.WithTransport(transport)) - h.clients[nodeID] = c - return c + h.clients[nodeID] = &nodeFileClient{client: c, transport: transport} + return c, nil +} + +// do sends a request to the HTTP server of nodeID. +func (h *HTTPFileStager) do(nodeID string, req *http.Request) (*http.Response, error) { + client, err := h.clientFor(nodeID) + if err != nil { + return nil, err + } + return client.Do(req) +} + +// ForgetNode drops the client of a node that left, and closes the idle streams +// that its transport holds. Without it the map grows by one client for every +// worker that ever connected, and a stream of a tunnel that is gone stays open. +// It is safe on a nil receiver. +func (h *HTTPFileStager) ForgetNode(nodeID string) { + if h == nil { + return + } + h.clientsMu.Lock() + entry, ok := h.clients[nodeID] + delete(h.clients, nodeID) + h.clientsMu.Unlock() + if ok { + entry.transport.CloseIdleConnections() + } } // ReleaseRemote removes one exact ephemeral key from a backend node. @@ -115,7 +159,7 @@ func (h *HTTPFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) if h.token != "" { req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.clientFor(nodeID).Do(req) + resp, err := h.do(nodeID, req) if err != nil { return fmt.Errorf("releasing %q from node %s: %w", key, nodeID, err) } @@ -153,7 +197,7 @@ func (h *HTTPFileStager) ReleaseRemoteRequest(ctx context.Context, nodeID, reque if h.token != "" { req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.clientFor(nodeID).Do(req) + resp, err := h.do(nodeID, req) if err != nil { return fmt.Errorf("releasing request inputs from node %s: %w", nodeID, err) } @@ -186,7 +230,10 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke return "", fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err) } // Fetched once per call so every retry reuses the same connection pool. - client := h.clientFor(nodeID) + client, err := h.clientFor(nodeID) + if err != nil { + return "", err + } // Probe: check if the remote already has the file with matching content hash. if remotePath, ok, probeErr := h.probeExisting(ctx, client, addr, localPath, key); probeErr != nil { @@ -821,7 +868,7 @@ func (h *HTTPFileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, loca req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.clientFor(nodeID).Do(req) + resp, err := h.do(nodeID, req) if err != nil { return fmt.Errorf("downloading from node %s: %w", nodeID, err) } @@ -877,7 +924,7 @@ func (h *HTTPFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (st req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.clientFor(nodeID).Do(req) + resp, err := h.do(nodeID, req) if err != nil { return "", fmt.Errorf("allocating temp file on node %s: %w", nodeID, err) } @@ -918,7 +965,7 @@ func (h *HTTPFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix st req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.clientFor(nodeID).Do(req) + resp, err := h.do(nodeID, req) if err != nil { return nil, fmt.Errorf("listing dir on node %s: %w", nodeID, err) } diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 5719dd48834b..9455f5c5cbc2 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -280,6 +280,32 @@ func NewTokenClientFactory(token string) BackendClientFactory { return &tokenClientFactory{token: token} } +// BackendDialerFor returns the dial function that reaches the backend processes +// of one worker, in the shape grpc.WithContextDialer takes. +type BackendDialerFor func(nodeID string) func(ctx context.Context, addr string) (net.Conn, error) + +// dialerClientFactory builds clients that reach their backend through a dialer +// chosen by node. +type dialerClientFactory struct { + token string + dialFor BackendDialerFor +} + +// NewDialerClientFactory returns a BackendClientFactory whose clients do not dial +// the address they are handed. The dialer of the node opens the connection, and +// the address becomes the name of a backend process of the worker. A tunnel +// carrier uses it. +// +// The clients are lazy, as every client of pkg/grpc is: a dial that fails shows +// up on the first call, and grpc.LastDialErrorOf reads why. +func NewDialerClientFactory(token string, dialFor BackendDialerFor) BackendClientFactory { + return &dialerClientFactory{token: token, dialFor: dialFor} +} + +func (f *dialerClientFactory) NewClient(nodeID, address string, parallel bool) grpc.Backend { + return grpc.NewClientWithDialer(address, parallel, nil, false, f.token, f.dialFor(nodeID)) +} + // WorkerNetDialerFor returns the dial function that reaches one worker's own // HTTP server, in the shape http.Transport.DialContext and // websocket.Dialer.NetDialContext take. It is keyed by node id, not address, From 6d788fb2a301d3ff07be8b3e0b210b35b33b041d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:12:50 +0000 Subject: [PATCH 064/162] feat(carrier): build the set of seams of the tunnel carrier NewTunnelSet builds the set that reaches workers through their tunnels: the control verbs over the HTTP control plane, the file stager over the bulk lane with no fallback, the gRPC clients, and the dialer of the HTTP server of a worker. The fan-out comes from the pgbus carrier. The work queue and the agent RPC are given by the caller until the parts that implement them exist, and the set is complete or it is an error. One function maps the errors of the tunnel onto the conditions that the rest of the frontend acts on. A route that does not exist is nodes.ErrNoRoute, with the error of the tunnel below it. A refusal of the worker about a backend keeps its identity and is marked as the answer of a host, so that code that decides whether a backend is dead can tell it from a failure of the transport. A peer that did not answer, the budget of the caller and a bulk lane that is down stay plain errors. The specs run the set through the holders on two replicas over a real PostgreSQL and real sockets: a load through the owner and through a replica that relays, a load that fails, a load that hangs and is cancelled, a worker that reconnects to the other replica, a file staged over the bulk lane, and a probe that stays fast while a large transfer runs. A transfer is refused when the worker has no bulk session. The dialer passes the same specs on the direct and the tunnel carrier: half-close, deadlines, and a dial that finds nothing. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../services/carrier/dial_conformance_test.go | 201 ++++++ core/services/carrier/tunnel.go | 164 +++++ .../carrier/tunnel_integration_test.go | 576 ++++++++++++++++++ 3 files changed, 941 insertions(+) create mode 100644 core/services/carrier/dial_conformance_test.go create mode 100644 core/services/carrier/tunnel.go create mode 100644 core/services/carrier/tunnel_integration_test.go diff --git a/core/services/carrier/dial_conformance_test.go b/core/services/carrier/dial_conformance_test.go new file mode 100644 index 000000000000..1a0fd077183a --- /dev/null +++ b/core/services/carrier/dial_conformance_test.go @@ -0,0 +1,201 @@ +package carrier_test + +import ( + "context" + "errors" + "io" + "net" + "time" + + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// dialHarness gives one carrier the same service to dial: a TCP server that +// reads its connection until the other side stops writing, and then answers with +// what it read and closes. +type dialHarness struct { + // dial is the dial function of the carrier for the node under test. + dial func(ctx context.Context) (net.Conn, error) + // unrouted is the same dial for a node that no worker holds. + unrouted func(ctx context.Context) (net.Conn, error) + // routesMissing says whether the carrier knows a node that is not connected. + routesMissing bool +} + +// Every carrier of WorkerNetDialerFor must pass these. The dialer serves the +// file stager, the control client and the backend-logs proxy, and each of them +// relies on the same behaviour of the connection that comes back. +var _ = Describe("WorkerNetDialerFor conformance", func() { + // echoService answers "got:" after the end of the request. + echoService := func() net.Listener { + GinkgoHelper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = lis.Close() }) + go func() { + for { + c, err := lis.Accept() + if err != nil { + return + } + go func() { + defer func() { _ = c.Close() }() + data, _ := io.ReadAll(c) + _, _ = c.Write([]byte("got:" + string(data))) + }() + } + }() + return lis + } + + harnesses := map[string]func() dialHarness{ + "the direct dialer of the NATS carrier": func() dialHarness { + lis := echoService() + set, err := carrier.NewNATSSet(carrier.NATSOptions{ + Client: &closingBus{FakeBus: testutil.NewFakeBus()}, Registry: noModels{}, + HTTPAddrFor: func(string) (string, error) { return lis.Addr().String(), nil }, + }) + Expect(err).ToNot(HaveOccurred()) + dial := set.Dialer("n") + return dialHarness{ + dial: func(ctx context.Context) (net.Conn, error) { return dial(ctx, "tcp", lis.Addr().String()) }, + unrouted: func(ctx context.Context) (net.Conn, error) { return dial(ctx, "tcp", "127.0.0.1:1") }, + } + }, + "the tunnel dialer on the replica that holds the tunnel": func() dialHarness { + fx := newDialFixture(echoService()) + return dialHarness{dial: fx.dialVia(fx.owner), unrouted: fx.unroutedVia(fx.owner), routesMissing: true} + }, + "the tunnel dialer through the replica that does not": func() dialHarness { + fx := newDialFixture(echoService()) + return dialHarness{dial: fx.dialVia(fx.other), unrouted: fx.unroutedVia(fx.other), routesMissing: true} + }, + } + + for name, build := range harnesses { + Describe(name, func() { + var h dialHarness + BeforeEach(func() { h = build() }) + + It("carries the end of the request to the service, and the whole answer back", func() { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + conn, err := h.dial(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + _, err = conn.Write([]byte("request")) + Expect(err).ToNot(HaveOccurred()) + closer, ok := conn.(interface{ CloseWrite() error }) + Expect(ok).To(BeTrue(), "the connection must half-close, or a request that ends with EOF waits for ever") + Expect(closer.CloseWrite()).To(Succeed()) + + Expect(conn.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + got, err := io.ReadAll(conn) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("got:request")) + }) + + It("honours the deadline that the caller puts on the connection", func() { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + conn, err := h.dial(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + // The service answers after the request ends, and this one does not + // end it. + Expect(conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))).To(Succeed()) + begun := time.Now() + _, err = conn.Read(make([]byte, 1)) + var netErr net.Error + Expect(errors.As(err, &netErr)).To(BeTrue(), "%v", err) + Expect(netErr.Timeout()).To(BeTrue()) + Expect(time.Since(begun)).To(BeNumerically("<", 5*time.Second)) + }) + + It("leaves no deadline armed on the connection it hands back", func() { + short, cancel := context.WithTimeout(context.Background(), 400*time.Millisecond) + defer cancel() + conn, err := h.dial(short) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + time.Sleep(700 * time.Millisecond) + _, err = conn.Write([]byte("later")) + Expect(err).ToNot(HaveOccurred(), "the budget of the dial is not the budget of the conversation") + }) + + It("fails a dial whose context is already over, without blaming a route", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := h.dial(ctx) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeFalse(), "the caller gave up; no route was tried") + }) + + It("fails a dial that finds nothing", func() { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, err := h.unrouted(ctx) + Expect(err).To(HaveOccurred()) + if h.routesMissing { + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeTrue(), "a node that no worker holds has no route: %v", err) + } + }) + }) + } +}) + +// dialFixture is two replicas, one of which holds the tunnel of a worker whose +// http service is the echo service. +type dialFixture struct { + owner, other *replica + nodeID string +} + +func newDialFixture(service net.Listener) *dialFixture { + GinkgoHelper() + ctx := context.Background() + db, dsn := testutil.SetupTestDBWithDSN() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + return buildDialFixture(ctx, db, dsn, service) +} + +func buildDialFixture(ctx context.Context, db *gorm.DB, dsn string, service net.Listener) *dialFixture { + GinkgoHelper() + fx := &dialFixture{} + clusterR := cluster.NewRegistry(db) + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken)} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + fx.nodeID = node.ID + Expect(nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(workerToken))).To(Succeed()) + + fx.owner = startReplica(ctx, "replica-a", db, dsn, clusterR, nodeReg) + fx.other = startReplica(ctx, "replica-b", db, dsn, clusterR, nodeReg) + + stop := dialInferenceOnly(fx.owner.url, fx.nodeID, workerToken, service.Addr().String()) + DeferCleanup(stop) + Eventually(func() bool { return fx.owner.tunnels.Holds(fx.nodeID) }, "15s", "50ms").Should(BeTrue()) + return fx +} + +func (fx *dialFixture) dialVia(r *replica) func(context.Context) (net.Conn, error) { + dial := r.set.Dialer(fx.nodeID) + return func(ctx context.Context) (net.Conn, error) { return dial(ctx, "tcp", fx.nodeID+".worker.invalid:80") } +} + +func (fx *dialFixture) unroutedVia(r *replica) func(context.Context) (net.Conn, error) { + dial := r.set.Dialer("a-node-that-no-worker-holds") + return func(ctx context.Context) (net.Conn, error) { return dial(ctx, "tcp", "nowhere.worker.invalid:80") } +} diff --git a/core/services/carrier/tunnel.go b/core/services/carrier/tunnel.go new file mode 100644 index 000000000000..d7499063c587 --- /dev/null +++ b/core/services/carrier/tunnel.go @@ -0,0 +1,164 @@ +package carrier + +import ( + "context" + "errors" + "fmt" + "net" + "time" + + mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/tunnel" +) + +// TunnelOptions is what NewTunnelSet needs. +// +// The set has seven members. Four of them belong to this carrier and are built +// here: the control verbs, the file stager, the client factory and the dialer of +// the HTTP server of a worker. The fan-out comes from NewPgbusFanout. The queue +// and the agent RPC of the tunnel carrier are provided by the caller, because the +// parts that implement them are separate from the parts that reach a worker. +type TunnelOptions struct { + // Epoch is the epoch of the cluster row the set is built for. + Epoch int64 + + // Dialer opens streams to workers, through the replica that holds the tunnel + // of each one. + Dialer *tunnel.WorkerDialer + + // Fanout carries the broadcasts. + Fanout *Fanout + // WorkQueue and Agents are the other two members that the set needs. + WorkQueue messaging.WorkQueue + Agents mcpTools.AgentControl + + // Registry finds the nodes that hold a model, for the verbs that fan out. + Registry nodes.ModelLocator + // InstallTimeout and UpgradeTimeout bound the install and upgrade verbs. + InstallTimeout, UpgradeTimeout time.Duration + // Token is the registration token. A worker checks it on its control plane, + // its file routes and its backend processes. + Token string +} + +// TunnelSet is a Set for the tunnel carrier, with the pieces that hold a cache +// per node. Forget drops them when a node leaves. +type TunnelSet struct { + *Set + control *nodes.ControlClient + files *nodes.HTTPFileStager +} + +// ForgetNode drops the clients that this set caches for a node and closes the +// idle streams they hold. Call it when a node is removed or loses its tunnel for +// good. +func (t *TunnelSet) ForgetNode(nodeID string) { + t.control.ForgetNode(nodeID) + t.files.ForgetNode(nodeID) +} + +// NewTunnelSet builds the set of seam implementations that reach workers through +// their tunnels. The set is complete or it is an error, as for NATS. +func NewTunnelSet(o TunnelOptions) (*TunnelSet, error) { + if o.Dialer == nil { + return nil, errors.New("tunnel set needs a dialer") + } + if o.Fanout == nil { + return nil, errors.New("tunnel set needs a fan-out") + } + + // The control verbs, the logs and the health of a worker go over the + // inference lane. The file transfer goes over the bulk lane and never falls + // back to the other one. + httpDial := httpDialerFor(o.Dialer) + bulkDial := httpDialerFor(o.Dialer, tunnel.WithBulkLane()) + + control := nodes.NewControlClient(httpDial, o.Token) + files := nodes.NewHTTPFileStager(func(nodeID string) (string, error) { + // The host is never dialled: the dialer opens a stream on the tunnel of + // the node. It names the node in a log line and in the URL. + return nodes.WorkerHTTPHost(nodeID, ""), nil + }, o.Token, bulkDial) + + set := &Set{ + Name: cluster.CarrierTunnel, + Epoch: o.Epoch, + Broadcaster: o.Fanout.Broadcaster, + OnReconnect: o.Fanout.OnReconnect, + WorkQueue: o.WorkQueue, + Commands: nodes.NewTunnelControl(o.Registry, control, o.InstallTimeout, o.UpgradeTimeout), + Files: files, + Clients: nodes.NewDialerClientFactory(o.Token, grpcDialerFor(o.Dialer)), + Dialer: httpDial, + Agents: o.Agents, + Close: o.Fanout.Close, + } + if err := set.Validate(); err != nil { + return nil, err + } + return &TunnelSet{Set: set, control: control, files: files}, nil +} + +// httpDialerFor returns the dialer of the HTTP server of a worker. +func httpDialerFor(d *tunnel.WorkerDialer, opts ...tunnel.DialOption) nodes.WorkerNetDialerFor { + return func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + dial := d.DialerFor(nodeID, tunnel.StreamTagHTTP, opts...) + return func(ctx context.Context, network, addr string) (net.Conn, error) { + conn, err := dial(ctx, network, addr) + return conn, mapDialError(err) + } + } +} + +// grpcDialerFor returns the dialer of the backend processes of a worker. +func grpcDialerFor(d *tunnel.WorkerDialer) nodes.BackendDialerFor { + return func(nodeID string) func(context.Context, string) (net.Conn, error) { + dial := d.GRPCDialerFor(nodeID) + return func(ctx context.Context, addr string) (net.Conn, error) { + conn, err := dial(ctx, addr) + return conn, mapDialError(err) + } + } +} + +// mapDialError is the one place where the errors of the tunnel become the +// conditions that the rest of the frontend acts on. It keeps the chain of the +// cause, so that the specific condition stays matchable below the one that is +// acted on. +// +// - A route that does not exist is nodes.ErrNoRoute, the one routing sentinel +// that the scheduler, the health monitor and the upgrade fallback match on. +// The sentinel of the tunnel package stays below it. +// - A worker's own refusal that is evidence about a backend keeps its identity +// and carries nodes.ErrNoRoute never. It is also marked as the answer of a +// host about a backend (grpc.BackendAnswer), so the code that decides whether +// a backend is dead can tell it from a failure of the transport. +// - A peer that did not answer, the budget of the caller and a bulk lane that +// is down are none of those. They stay plain errors. +func mapDialError(err error) error { + switch { + case err == nil: + return nil + case tunnel.IsWorkerAnswer(err): + return &workerAnswerError{err: err} + case errors.Is(err, tunnel.ErrNoRoute): + return fmt.Errorf("%w: %w", nodes.ErrNoRoute, err) + } + return err +} + +// workerAnswerError is a refusal that the worker wrote about a backend. It is a +// grpc.BackendAnswer. +type workerAnswerError struct{ err error } + +func (e *workerAnswerError) Error() string { return e.err.Error() } + +func (e *workerAnswerError) Unwrap() error { return e.err } + +// IsBackendAnswer marks the error as evidence about the backend: the worker +// answered, and the backend it named is not there, like a refused connection on a +// direct dial. +func (e *workerAnswerError) IsBackendAnswer() bool { return true } diff --git a/core/services/carrier/tunnel_integration_test.go b/core/services/carrier/tunnel_integration_test.go new file mode 100644 index 000000000000..6041aa649a77 --- /dev/null +++ b/core/services/carrier/tunnel_integration_test.go @@ -0,0 +1,576 @@ +package carrier_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "slices" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/labstack/echo/v4" + ggrpc "google.golang.org/grpc" + "gorm.io/gorm" + + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/tunnel/slowlink" + "github.com/mudler/LocalAI/core/services/worker" + "github.com/mudler/LocalAI/core/services/workerctl" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fakeBackend is a gRPC backend that a spec can make load, fail to load, or hang +// while it loads. +type fakeBackend struct { + pb.UnimplementedBackendServer + + mu sync.Mutex + loads []string + cancelled []string + loading chan string +} + +func (b *fakeBackend) Health(context.Context, *pb.HealthMessage) (*pb.Reply, error) { + return &pb.Reply{Message: []byte("OK")}, nil +} + +func (b *fakeBackend) LoadModel(ctx context.Context, in *pb.ModelOptions) (*pb.Result, error) { + b.mu.Lock() + b.loads = append(b.loads, in.Model) + b.mu.Unlock() + switch in.Model { + case "fails": + return &pb.Result{Success: false, Message: "out of memory"}, nil + case "hangs": + select { + case b.loading <- in.Model: + default: + } + <-ctx.Done() + b.mu.Lock() + b.cancelled = append(b.cancelled, in.Model) + b.mu.Unlock() + return nil, ctx.Err() + } + return &pb.Result{Success: true, Message: "loaded " + in.Model}, nil +} + +func (b *fakeBackend) Predict(_ context.Context, in *pb.PredictOptions) (*pb.Reply, error) { + return &pb.Reply{Message: []byte("echo: " + in.Prompt)}, nil +} + +func (b *fakeBackend) cancelledModels() []string { + b.mu.Lock() + defer b.mu.Unlock() + return slices.Clone(b.cancelled) +} + +// controlPlane is the control plane of the fake worker. It answers the verbs of a +// load and counts what it saw. +type controlPlane struct { + mu sync.Mutex + address string + installs []workerctl.BackendInstallRequest + stops []workerctl.ModelStopRequest + renewals int +} + +func (c *controlPlane) handler() http.Handler { + mux := http.NewServeMux() + body := func(r *http.Request, into any) { + raw, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(raw, into) + } + mux.HandleFunc(workerctl.PathOf(workerctl.VerbBackendInstall), func(w http.ResponseWriter, r *http.Request) { + var req workerctl.BackendInstallRequest + body(r, &req) + c.mu.Lock() + c.installs = append(c.installs, req) + address := c.address + c.mu.Unlock() + w.Header().Set("Content-Type", workerctl.ContentTypeStream) + ev, _ := json.Marshal(workerctl.BackendInstallProgressEvent{OpID: req.OpID, Percentage: 50}) + _, _ = fmt.Fprintf(w, "{\"progress\":%s}\n", ev) + reply, _ := json.Marshal(workerctl.BackendInstallReply{Success: true, Address: address, ProcessInstance: "i1", ReportsOperations: true}) + _, _ = fmt.Fprintf(w, "{\"reply\":%s}\n", reply) + }) + mux.HandleFunc(workerctl.PathOf(workerctl.VerbModelStop), func(w http.ResponseWriter, r *http.Request) { + var req workerctl.ModelStopRequest + body(r, &req) + c.mu.Lock() + c.stops = append(c.stops, req) + c.mu.Unlock() + // A repeat answers like the first: Terminated, never an error. + _ = json.NewEncoder(w).Encode(workerctl.ModelStopReply{Matched: true, Terminated: true}) + }) + mux.HandleFunc(workerctl.PathOf(workerctl.VerbModelOp), func(w http.ResponseWriter, r *http.Request) { + var req workerctl.OperationRequest + body(r, &req) + c.mu.Lock() + c.renewals++ + c.mu.Unlock() + _ = json.NewEncoder(w).Encode(workerctl.OperationReply{Renewed: req.Renew, Completed: req.Complete}) + }) + mux.HandleFunc(workerctl.PathOf(workerctl.VerbBackendList), func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"backends":[]}`) + }) + mux.HandleFunc(workerctl.PathOf(workerctl.VerbModelsRunning), func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"models":[]}`) + }) + return mux +} + +func (c *controlPlane) seenStops() []workerctl.ModelStopRequest { + c.mu.Lock() + defer c.mu.Unlock() + return slices.Clone(c.stops) +} + +func (c *controlPlane) seenRenewals() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.renewals +} + +// replica is one frontend: its own tunnel registry, peer link, relay and carrier +// set, behind its own HTTP server, over the database that all replicas share. +type replica struct { + id string + url string + tunnels *tunnel.Registry + sessions *tunnel.PeerSessions + pool *tunnel.PeerPool + set *carrier.TunnelSet + active *atomic.Pointer[carrier.Set] + + commands *carrier.Commands + files *carrier.Files + clients *carrier.Clients +} + +// startReplica starts one frontend: its own tunnel registry, peer link, relay and +// carrier set, behind its own HTTP server, over the database that all replicas +// share. +func startReplica(ctx context.Context, id string, db *gorm.DB, dsn string, clusterR *cluster.Registry, nodeReg *nodes.NodeRegistry) *replica { + GinkgoHelper() + r := &replica{id: id} + cred := cluster.NewPeerCredential() + r.tunnels = tunnel.NewRegistry(clusterR, id) + r.sessions = tunnel.NewPeerSessions(tunnel.NewRelay(r.tunnels).Stream) + DeferCleanup(r.sessions.Close) + r.pool = tunnel.NewPeerPool(id, cred, clusterR) + DeferCleanup(r.pool.Close) + + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(nodeReg, r.tunnels)) + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(clusterR, r.sessions.Accept)) + srv := httptest.NewServer(e) + DeferCleanup(srv.Close) + r.url = srv.URL + Expect(clusterR.RegisterPeer(ctx, id, "test", 0, "", strings.TrimPrefix(srv.URL, "http://"), cred.Hash())).To(Succeed()) + + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(fan.Close) + fake := newFakeCarrier(cluster.CarrierTunnel, 1) + set, err := carrier.NewTunnelSet(carrier.TunnelOptions{ + Epoch: 1, + Dialer: tunnel.NewWorkerDialer(r.tunnels, r.pool), + Fanout: fan, + WorkQueue: fake.queue, + Agents: fake.agents, + Registry: nodeReg, + InstallTimeout: 30 * time.Second, + UpgradeTimeout: 30 * time.Second, + Token: registrationToken, + }) + Expect(err).ToNot(HaveOccurred()) + r.set = set + r.active = &atomic.Pointer[carrier.Set]{} + r.active.Store(set.Set) + r.commands = carrier.NewCommands(r.active) + r.files = carrier.NewFiles(r.active) + r.clients = carrier.NewClients(r.active) + return r +} + +const ( + registrationToken = "shared-registration-token" + workerToken = "own-tunnel-credential-of-w1" +) + +func hashOf(secret string) string { + sum := sha256.Sum256([]byte(secret)) + return hex.EncodeToString(sum[:]) +} + +var _ = Describe("The tunnel carrier through the holders, end to end", func() { + var ( + ctx context.Context + db *gorm.DB + dsn string + nodeID string + + clusterR *cluster.Registry + nodeReg *nodes.NodeRegistry + a, b *replica + + backend *fakeBackend + control *controlPlane + grpcAddr string + + httpAddr string + stagingDir string + ) + + newReplica := func(id string) *replica { + GinkgoHelper() + return startReplica(ctx, id, db, dsn, clusterR, nodeReg) + } + + // attach starts a worker tunnel to url. Its services reach the fake backend + // and the HTTP server of the fake worker. + attach := func(url string) *worker.Tunnel { + GinkgoHelper() + t, err := worker.StartTunnel(ctx, worker.TunnelConfig{ + FrontendURL: url, + NodeID: nodeID, + Token: func() string { return workerToken }, + Services: map[string]worker.LocalService{ + tunnel.StreamTagGRPC: func(ctx context.Context, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", grpcAddr) + }, + tunnel.StreamTagHTTP: func(ctx context.Context, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", httpAddr) + }, + }, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(t.Close) + return t + } + + ownerIs := func(r *replica) func() bool { + return func() bool { + owner, _, err := clusterR.Owner(ctx, nodeID) + return err == nil && owner == r.id + } + } + + BeforeEach(func() { + ctx = context.Background() + db, dsn = testutil.SetupTestDBWithDSN() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + var err error + nodeReg, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken)} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + nodeID = node.ID + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hashOf(workerToken))).To(Succeed()) + + // The fake worker: a gRPC backend and an HTTP server that serves the + // control plane and the file transfer of the real worker. + backend = &fakeBackend{loading: make(chan string, 1)} + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + srv := ggrpc.NewServer() + pb.RegisterBackendServer(srv, backend) + go func() { _ = srv.Serve(lis) }() + DeferCleanup(srv.Stop) + grpcAddr = lis.Addr().String() + + control = &controlPlane{address: grpcAddr} + stagingDir = GinkgoT().TempDir() + hl, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + httpAddr = hl.Addr().String() + Expect(hl.Close()).To(Succeed()) + httpSrv, err := nodes.StartFileTransferServerWithControl(httpAddr, stagingDir, stagingDir, stagingDir, registrationToken, 1<<30, nil, nil, control.handler()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { nodes.ShutdownFileTransferServer(httpSrv) }) + + a, b = newReplica("replica-a"), newReplica("replica-b") + }) + + It("loads a model through the replica that holds the tunnel", func() { + attach(a.url) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + + var progress []float64 + reply, err := a.commands.InstallBackendOp(nodeID, "llama-cpp", "m", "", 0, "op-1", "load-1", time.Minute, func(ev workerctl.BackendInstallProgressEvent) { + progress = append(progress, ev.Percentage) + }) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + Expect(reply.Address).To(Equal(grpcAddr)) + Expect(progress).To(Equal([]float64{50})) + + client := a.clients.NewClient(nodeID, reply.Address, false) + res, err := client.LoadModel(ctx, &pb.ModelOptions{Model: "ok"}) + Expect(err).ToNot(HaveOccurred()) + Expect(res.Success).To(BeTrue()) + out, err := client.Predict(ctx, &pb.PredictOptions{Prompt: "hi"}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out.Message)).To(Equal("echo: hi")) + + renewed, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"load-1"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(renewed.Renewed).To(ConsistOf("load-1")) + Expect(a.commands.PingNode(nodeID)).To(Succeed()) + }) + + It("loads a model through a replica that does not hold the tunnel, over the peer link", func() { + attach(a.url) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + + reply, err := b.commands.InstallBackendOp(nodeID, "llama-cpp", "m", "", 0, "op-2", "load-2", time.Minute, nil) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + + client := b.clients.NewClient(nodeID, reply.Address, false) + res, err := client.LoadModel(ctx, &pb.ModelOptions{Model: "ok"}) + Expect(err).ToNot(HaveOccurred()) + Expect(res.Success).To(BeTrue()) + Expect(a.sessions.Holds("replica-b")).To(BeTrue(), "the request crossed a link from the replica that does not hold the tunnel to the one that does") + Expect(b.tunnels.Holds(nodeID)).To(BeFalse()) + }) + + It("reports a failing load as the answer of the backend, and never as a missing route", func() { + attach(a.url) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + + for _, r := range []*replica{a, b} { + client := r.clients.NewClient(nodeID, grpcAddr, false) + res, err := client.LoadModel(ctx, &pb.ModelOptions{Model: "fails"}) + Expect(err).ToNot(HaveOccurred(), r.id) + Expect(res.Success).To(BeFalse(), r.id) + Expect(res.Message).To(Equal("out of memory"), r.id) + } + }) + + It("cancels a load that hangs, and stops its operation twice without an error", func() { + attach(a.url) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + + loadCtx, cancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { + client := b.clients.NewClient(nodeID, grpcAddr, false) + _, err := client.LoadModel(loadCtx, &pb.ModelOptions{Model: "hangs"}) + done <- err + }() + Eventually(backend.loading, "10s").Should(Receive(Equal("hangs"))) + + cancel() + var err error + Eventually(done, "10s").Should(Receive(&err)) + Expect(err).To(HaveOccurred()) + Eventually(backend.cancelledModels, "10s").Should(ConsistOf("hangs"), "the cancellation must reach the backend of the worker through the relay and the tunnel") + + for range 2 { + stop, err := b.commands.StopLoadOperation(ctx, nodeID, workerctl.ModelStopRequest{ + ModelName: "m", ProcessKey: "m#0", ExpectedAddress: grpcAddr, OperationID: "load-3", Force: true}) + Expect(err).ToNot(HaveOccurred()) + Expect(stop.Terminated).To(BeTrue()) + } + Expect(control.seenStops()).To(HaveLen(2)) + }) + + It("keeps renewing across a worker that reconnects to the other replica", func() { + first := attach(a.url) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"load-4"}}) + Expect(err).ToNot(HaveOccurred()) + + // The worker loses replica A and dials replica B. + Expect(first.Close()).To(Succeed()) + attach(b.url) + Eventually(ownerIs(b), "20s", "50ms").Should(BeTrue()) + + // Replica A holds nothing now, and reaches the worker over the link. + Eventually(func() error { + _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"load-4"}}) + return err + }, "20s", "100ms").Should(Succeed()) + Expect(a.tunnels.Holds(nodeID)).To(BeFalse()) + Expect(b.sessions.Holds("replica-a")).To(BeTrue()) + Expect(control.seenRenewals()).To(BeNumerically(">=", 2)) + }) + + It("reports a worker that is connected nowhere as ErrNoRoute, within the cadence of a renewal", func() { + begun := time.Now() + _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse(), "absence must not reach the caller") + Expect(time.Since(begun)).To(BeNumerically("<", 5*time.Second)) + Expect(errors.Is(a.commands.PingNode(nodeID), nodes.ErrNoRoute)).To(BeTrue()) + }) + + It("stages a file over the bulk lane, from either replica", func() { + attach(a.url) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + Eventually(func() bool { return a.tunnels.Holds(nodeID) }, "10s").Should(BeTrue()) + + src := filepath.Join(GinkgoT().TempDir(), "weights.bin") + payload := make([]byte, 6<<20) + for i := range payload { + payload[i] = byte(i * 7) + } + Expect(os.WriteFile(src, payload, 0o600)).To(Succeed()) + + for _, r := range []*replica{a, b} { + var remote string + Eventually(func() error { + var err error + remote, err = r.files.EnsureRemote(ctx, nodeID, src, "models/"+r.id+"/weights.bin") + return err + }, "30s", "200ms").Should(Succeed(), r.id) + got, err := os.ReadFile(remote) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal(payload), r.id) + } + }) + + It("keeps a small call fast while a large transfer runs on the same worker", func() { + // The worker reaches the frontend through a link with a rate and a delay, + // as a worker in another network does. + link, err := slowlink.New(strings.TrimPrefix(a.url, "http://"), 8e6, 10*time.Millisecond) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(link.Close) + attach("http://" + link.Addr()) + Eventually(ownerIs(a), "20s", "50ms").Should(BeTrue()) + Eventually(func() bool { return a.tunnels.Holds(nodeID) }, "10s").Should(BeTrue()) + // The bulk lane joins a moment after the inference lane. + time.Sleep(time.Second) + + src := filepath.Join(GinkgoT().TempDir(), "big.bin") + Expect(os.WriteFile(src, make([]byte, 16<<20), 0o600)).To(Succeed()) + + var ( + samples []time.Duration + stop = make(chan struct{}) + wg sync.WaitGroup + ) + wg.Go(func() { + for { + select { + case <-stop: + return + case <-time.After(10 * time.Millisecond): + } + begun := time.Now() + if _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"op"}}); err != nil { + return + } + samples = append(samples, time.Since(begun)) + } + }) + begun := time.Now() + _, err = a.files.EnsureRemote(ctx, nodeID, src, "models/big.bin") + elapsed := time.Since(begun) + close(stop) + wg.Wait() + Expect(err).ToNot(HaveOccurred()) + + Expect(elapsed).To(BeNumerically(">", 500*time.Millisecond), "the transfer must last long enough to be measured") + Expect(len(samples)).To(BeNumerically(">=", 10), "probes completed during the transfer") + slices.Sort(samples) + p99 := samples[min(len(samples)-1, len(samples)*99/100)] + AddReportEntry("probe during a bulk transfer", map[string]any{ + "transfer s": elapsed.Seconds(), "probes": len(samples), "p50 ms": samples[len(samples)/2].Milliseconds(), "p99 ms": p99.Milliseconds(), + }) + Expect(p99).To(BeNumerically("<", 100*time.Millisecond), + "a probe on the inference lane must not wait behind the bytes of a transfer on the bulk lane") + }) + + It("refuses a transfer when the worker has no bulk session, and sends nothing on the other lane", func() { + // A worker that dials the inference lane only: the case of a worker that + // predates the bulk lane, or one whose bulk session is down. + inferenceOnly := dialInferenceOnly(a.url, nodeID, workerToken, httpAddr) + DeferCleanup(inferenceOnly) + Eventually(ownerIs(a), "15s", "50ms").Should(BeTrue()) + + src := filepath.Join(GinkgoT().TempDir(), "x.bin") + Expect(os.WriteFile(src, []byte("bytes"), 0o600)).To(Succeed()) + for _, r := range []*replica{a, b} { + _, err := r.files.EnsureRemote(ctx, nodeID, src, "models/x.bin") + Expect(err).To(HaveOccurred(), r.id) + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeFalse(), "%s: a lane that is down must not demote a worker whose model calls work: %v", r.id, err) + } + Expect(filepath.Join(stagingDir, "models", "x.bin")).ToNot(BeAnExistingFile()) + // The inference lane still carries a control call. + Expect(a.commands.PingNode(nodeID)).To(Succeed()) + }) + + It("carries broadcasts between the replicas through the holders", func() { + subject := messagingtest.BroadcastRootSubjects["state"] + var got atomic.Int32 + _, err := carrier.NewBroadcaster(b.active).Subscribe(subject, func([]byte) { got.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + Expect(carrier.NewBroadcaster(a.active).Publish(subject, map[string]string{"k": "v"})).To(Succeed()) + Eventually(got.Load, "10s").Should(Equal(int32(1))) + }) +}) + +// dialInferenceOnly connects a fake worker that holds the inference lane and no +// bulk lane. It serves the http tag from httpAddr. The returned function closes +// it. +func dialInferenceOnly(frontend, nodeID, token, httpAddr string) func() { + GinkgoHelper() + url := "ws" + strings.TrimPrefix(frontend, "http") + tunnel.ConnectPath + "?id=" + nodeID + header := http.Header{"Authorization": []string{"Bearer " + token}} + ws, _, err := tunnel.NewDialer(5*time.Second).Dial(url, header) + Expect(err).ToNot(HaveOccurred()) + sess, err := tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + go func() { + for { + st, err := sess.AcceptStream() + if err != nil { + return + } + go func() { + tag, _, err := tunnel.ReadStreamRequest(st) + if err != nil || tag != tunnel.StreamTagHTTP { + _ = tunnel.WriteStreamRefusal(st, tunnel.ErrStreamTagUnknown) + _ = st.Close() + return + } + local, err := net.Dial("tcp", httpAddr) + if err != nil { + _ = tunnel.WriteStreamRefusal(st, tunnel.ErrStreamTargetUnavailable) + _ = st.Close() + return + } + _ = tunnel.WriteStreamAccepted(st) + _ = tunnel.Splice(st, local) + }() + } + }() + return func() { _ = sess.Close() } +} From 655e1870888cb0aff478e86b6dd1d26c123c61d2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:13:42 +0000 Subject: [PATCH 065/162] docs(agents): describe the HTTP control plane, the dial seam and the peer link The contributor guide now says how the second carrier of the control verbs works (one implementation over a link, and the HTTP server of the worker), how a dial that fails in the transport is told from a dead backend, how a replica reaches a worker that another replica holds, and which of the open items are closed. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 139 ++++++++++++++++++++++++++--------- 1 file changed, 104 insertions(+), 35 deletions(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 823f38f6d360..bfc67be80c4f 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -18,8 +18,8 @@ files to that spec. |---|---|---|---| | Fan-out | `messaging.Broadcaster` | `core/services/messaging` | NATS client | | Queues | `messaging.WorkQueue` (producer), `messaging.WorkConsumer` (worker), keyed by `messaging.WorkKind` | `core/services/messaging` (`workqueue.go`) | `NewNATSWorkQueue`, `NewNATSWorkConsumer` | -| Control verbs, frontend half | `nodes.NodeCommandSender`, `nodes.FileStager` | `core/services/nodes` | `RemoteUnloaderAdapter`, `S3NATSFileStager` over NATS request/reply; `HTTPFileStager` over HTTP | -| Control verbs, worker half | unexported `controlServer` (`handle`, `handleWithProgress`) | `core/services/worker` | `natsControlServer` | +| Control verbs, frontend half | `nodes.NodeCommandSender`, `nodes.FileStager` | `core/services/nodes` | `RemoteUnloaderAdapter`, `S3NATSFileStager` over NATS request/reply; `HTTPFileStager` over HTTP; `TunnelControl` over the HTTP control plane of a worker | +| Control verbs, worker half | unexported `controlServer` (`handle`, `handleWithProgress`) | `core/services/worker` | `natsControlServer`, `httpControlServer` | | Agent RPC, frontend half | `AgentControl` (package `mcp`) | `core/http/endpoints/mcp` | `nodes.NATSAgentControl` | | Agent RPC, worker half | unexported `agentRPCServer` | `core/cli/agent_worker.go` | `nodes.NATSAgentRPCServer` | | Dial to a worker | `nodes.BackendClientFactory`, `nodes.ModelProber`, `nodes.WorkerNetDialerFor` | `core/services/nodes` | direct dial | @@ -102,6 +102,20 @@ of the sessions that a replica holds, and `Splice`. The connect endpoint is and stores only its SHA-256 (`BackendNode.TunnelTokenHash`). The registration token of the deployment never opens a tunnel. Every registration mints a new credential, so the client reads it at every dial. +- The frontend resets every stream that a worker opens (`MaxIncomingStreams` is + 0 on the server side of a session), because nothing reads such a stream and an + accepted one would hold its window of unread data for as long as the session + lives. The worker side keeps the default: it accepts the streams of the + frontend. +- The credential is checked when a worker dials, and never again. A node that is + deleted, or whose credential a registration replaces, loses the session that + this replica holds (`Registry.Disconnect`), and a replica that shuts down + closes its registry (`Registry.Close`) before it leaves the instances table. A + session that another replica holds ends when its worker dials again and is + refused. A worker whose dial is refused with 401 or 403 registers again + (`TunnelConfig.Reauthorize`), once for each wait of its backoff. +- A claim has its own bound (10 seconds), and a claim that was written for a + request that already ended is released. - A refusal of a stream is one of four. `tunnel.IsWorkerAnswer` is true for the three that are evidence about a backend. `ErrStreamNotServed` says that the worker learned nothing, and it must never count as evidence. @@ -142,15 +156,41 @@ runs every carrier in `loadOperationCarriers` against it. Whether a worker names operations is the worker's own report (`BackendInstallReply.ReportsOperations`), not a property of the carrier. -Worker half: each verb is a `controlVerb`. A handler is typed with `unary`, -`withProgress` or `noReply` and registered with `handle` (one request of the -verb at a time on NATS, panic not recovered) or `handleWithProgress` (a -goroutine per request, progress published on the install-progress subject). -An undecodable body is still answered with the verb's typed refusal. The -`undecodable` error a `controlHandler` returns is read only by tests today: the -NATS server drops it. It is a recorded exception to the no-dead-code rule, kept -as the hook a carrier that signals a malformed request out of band (HTTP 400) -needs. +Worker half: each verb is a `controlVerb`, named by the constants of +`core/services/workerctl` (`workerctl.AllVerbs` lists them). A handler is typed +with `unary`, `withProgress` or `noReply` and registered with `handle` (one +request of the verb at a time on NATS, panic not recovered) or +`handleWithProgress` (a goroutine per request, progress published on the +install-progress subject). An undecodable body is still answered with the verb's +typed refusal on NATS. + +`httpControlServer` is the second server of the same handlers. It mounts each +verb at `workerctl.PathOf(verb)` below `/v1/control/`, on the HTTP server that the +worker already runs (`nodes.StartFileTransferServerWithControl`, behind the same +bearer token as the file routes). It delivers requests of a verb concurrently, +where NATS delivers them one by one, so a handler must be safe to call from +several goroutines. The answer of a verb is 200 with the reply, and the refusal +of the worker stays in the reply. A body that cannot be read, including a body +that does not decode (the `undecodable` return of a `controlHandler`), is 400. A +verb the worker does not serve is 404 with a body that says so, and a method other +than POST is 405. A verb with progress streams lines (`workerctl.Envelope`): any +number of progress lines and then exactly one reply line. A body that ends +without a reply line is a link that broke, and it says nothing about the work. +A verb that only one server can serve is simply not registered on the other +(`files.*` need object storage on both). + +The frontend half is one implementation of the verbs over a `controlLink` +(`nodeControl` in `unloader.go`). The link holds how one request travels and how +its failure is told apart: `natsLink` for NATS and `httpLink` for the control +plane of a worker (`ControlClient` is the HTTP client). Adding a carrier of the +verbs is a new link, not a copy of the verbs. The link decides four things: what +is `ErrNoRoute`, whether a wait that ran out means that the worker may still be +installing (`waitExpired`), whether a missing acknowledgement of `backend.stop` +may be read as an older worker (`acknowledgementMissing`: true only for NATS), and +how the carrier is named in a log line. An HTTP link never reads silence as a +stop that was done. A 404 is a worker that does not serve the verb. It becomes +`ErrNoRoute` only for `backend.upgrade`, which keeps the fallback to the older +install, and it is a plain error for every other verb. ## Agent RPC @@ -169,12 +209,55 @@ client for SmartRouter, HealthMonitor and the reconciler's default `ModelProber`. The node id is there because a dialer that must know which node it reaches, as a tunnel does, cannot recover it from the address; the direct factory ignores it. `ModelProber.Probe(ctx, nodeID, address)` likewise. +`NewDialerClientFactory` builds clients through `grpc.NewClientWithDialer`: the +address is then the name of a backend process of the worker, and gRPC does not +resolve it. `WorkerNetDialerFor` returns the dial function for one worker's own HTTP server. `DirectWorkerNetDialer` dials the address it is handed. It serves -`HTTPFileStager` and the backend-logs proxy (HTTP and WebSocket). -`HTTPFileStager.clientFor` keeps one HTTP client per node, because the idle -pool is keyed by host and port and two workers can report the same address. +`HTTPFileStager`, `ControlClient` and the backend-logs proxy (HTTP and +WebSocket). `HTTPFileStager.clientFor` keeps one HTTP client per node, because +the idle pool is keyed by host and port and two workers can report the same +address, and `ForgetNode` drops one. A worker that holds a tunnel has no address. +`nodes.WorkerHTTPHost` gives it a name under `.invalid` for the URL, which is +never dialled, and the proxy of the environment is not used for such a host. + +A dial that fails in the transport says nothing about the backend. gRPC reports +it as `codes.Unavailable`, the code of a backend that died, so the client keeps +the error of its last dial and `grpc.TransportFailureOf(client)` returns it, +looking through decorators (every decorator implements `Unwrap`). It is nil when +the dial worked and when the host answered about the backend (an error that +implements `grpc.BackendAnswer`). Code that decides a backend is dead must ask +it before it acts: the probe of the reconciler (`ProbeUnknown`), the per-model +check of the health monitor, the warm path of the router (`probeUnknown`), the +eviction of a remote model and the check of a cached remote model. None of them +may reap a row, count a miss or shut a model down on a transport failure. + +## The peer link and the relay + +A worker holds one tunnel and it lands on one replica. `tunnel.WorkerDialer` is +the one door to a worker: the replica that holds the tunnel opens the stream, and +any other replica relays through the owner over a link that the pool of peers +(`tunnel.PeerPool`) holds and the owner accepts at `GET /api/cluster/peer` +(`tunnel.PeerSessions`, `tunnel.Relay`). A replica publishes its address and the +hash of its own peer credential in its instances row. The credential of the +replica is the only secret that opens a peer link; the registration token opens +none. A relayed stream carries two request frames, the relay frame (worker, lane, +remaining budget of the caller) and then the frame of the tunnel. The relay never +relays onward. + +The errors of a dial are kept apart. A route that does not exist carries +`tunnel.ErrNoRoute`, and the carrier maps it onto `nodes.ErrNoRoute` in one place +(`carrier.mapDialError`). Absence claims (`cluster.ErrNoConnection`, +`ErrInstanceNotFound`) never reach a caller. A refusal of the worker keeps its +identity and is no `ErrNoRoute`. A peer that does not answer is +`tunnel.ErrPeerUnreachable` and is not `ErrNoRoute`: a link between two frontends +says nothing about a worker, and the scheduler demotes a worker on `ErrNoRoute`. A +bulk lane that is down is `tunnel.ErrNoBulkSession` and is not `ErrNoRoute` +either. The budget of the caller is checked first and is never a route failure. + +A dial can ask for the bulk lane (`tunnel.WithBulkLane`) and then never falls +back to the lane of model calls. The file stager asks for it. ## Rules a carrier must keep @@ -264,28 +347,9 @@ other seams against a real server, also through Docker. ## Open items for a second carrier -- `DistributedModelStore.Range` (`core/services/nodes/distributed_store.go`) - builds a tokenless `model.Model` from `node.Address`. It dials nothing today, - but it is the one direct construction site left outside the dial seam. -- The worker's `files.ensure` handler passes the first caller's ctx into a - shared singleflight closure. The NATS server hands it `context.Background`. - A carrier with a request ctx needs `context.WithoutCancel` there, or one - cancelled caller fails the others. -- `HTTPFileStager` caches a client per node and never forgets one. There is no - `ForgetNode`. -- `clientFor` returns no error; a carrier with no dialer for a node needs that - path. - The agent worker still uses NATS directly for its connection and for agent events. `agents.NewEventBridge` takes a `Broadcaster`, so the frontend's bridge goes through the holder. -- The backend-logs proxy (`proxyHTTPToWorker`) starts from - `httpclient.HardenedTransport()`, which keeps - `Proxy: http.ProxyFromEnvironment`. With `HTTP_PROXY` set, a dialer that - routes by node would be handed the proxy address. Clear `Proxy` when the - dialer is not the direct one. -- `natsControlServer.subject` refuses a verb it has no subject for, so - registration fails. A verb that only one carrier serves needs a per-carrier - opt-out where the verbs are registered. - `WorkHandler` returns only an error. A carrier whose stream handler must send a terminal reply derives it from the `jobs..result` event the handler publishes on `events` (`handleMCPCIJob` does this). @@ -295,7 +359,12 @@ other seams against a real server, also through Docker. change `NATSDispatcher` and `EventBridge`, for example by binding `handleJob`'s publishes to `events` through a bridge view that shares the cancel registry. -- `controlHandler`'s `undecodable` return is the recorded exception described - under Control verbs. - Agent cancel has no production sender (`EventBridge.CancelExecution` has no caller), so it is not part of the agent RPC seam. +- The file stager over shared object storage (`S3NATSFileStager`) sends its + requests over NATS. The files verbs exist on the HTTP control plane, but no + frontend stager uses them yet. +- A worker that registers an address and holds a tunnel (a worker that can use + either carrier) is reached by the backend-logs proxy through the host that its + address names, and the proxy of the environment is used for it. Only a worker + without an address skips the proxy. From 20c48df41ad464525caf97731c2eab5399605c08 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:16:54 +0000 Subject: [PATCH 066/162] feat(application): forget a node in the active carrier when it loses its tunnel A node that an admin removes, or whose tunnel credential a registration replaces, now loses its session on this replica and the clients that the active carrier caches for it. A set offers ForgetNode for that. The set of the tunnel closes the control client and the stager of the node, and the NATS set closes the HTTP stager, which until now kept one client for each node that ever connected. The specs also pin that a replica publishes its peer address and the hash of its credential, and that it closes its tunnels before it leaves the instances table. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 13 +++ core/application/distributed_carrier_test.go | 88 +++++++++++++++++++ core/http/app.go | 4 +- .../services/carrier/dial_conformance_test.go | 18 ++++ core/services/carrier/nats.go | 5 ++ core/services/carrier/set.go | 5 ++ core/services/carrier/tunnel.go | 27 ++---- .../carrier/tunnel_integration_test.go | 11 ++- 8 files changed, 147 insertions(+), 24 deletions(-) diff --git a/core/application/distributed.go b/core/application/distributed.go index 6bb7d226d3d1..b2c675fdd621 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -101,6 +101,19 @@ type DistributedServices struct { shutdownOnce sync.Once } +// Disconnect ends the tunnel that this replica holds for a node, and drops what +// the active carrier caches for it. The node was removed or its credential was +// replaced. It reports whether this replica held a tunnel of the node. +func (ds *DistributedServices) Disconnect(nodeID string) bool { + held := ds.Tunnels.Disconnect(nodeID) + if ds.active != nil { + if set := ds.active.Load(); set != nil && set.ForgetNode != nil { + set.ForgetNode(nodeID) + } + } + return held +} + // Shutdown stops all distributed services in reverse initialization order. // It is safe to call on a nil receiver and is idempotent (uses sync.Once). func (ds *DistributedServices) Shutdown() { diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go index 13d7c651331d..c03fdde573af 100644 --- a/core/application/distributed_carrier_test.go +++ b/core/application/distributed_carrier_test.go @@ -3,13 +3,18 @@ package application import ( "context" "fmt" + "net/http" + "net/http/httptest" "os" "runtime" + "strings" "sync" + "time" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" "github.com/testcontainers/testcontainers-go" "github.com/testcontainers/testcontainers-go/wait" "gorm.io/gorm" @@ -146,6 +151,61 @@ var _ = Describe("distributed startup and the cluster carrier", func() { Expect(live()).To(ConsistOf("replica-b")) }) + It("publishes the address and the peer credential of the replica in its row", func() { + cfg.Distributed.PeerAddress = "10.0.0.7:8080" + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + inst, err := cluster.NewRegistry(db).Get(context.Background(), "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(inst.AdvertisedAddr).To(Equal("10.0.0.7:8080")) + Expect(inst.PeerTokenHash).To(HaveLen(64), "only the hash of the credential is published") + Expect(svc.PeerPool).ToNot(BeNil()) + Expect(svc.PeerSessions).ToNot(BeNil()) + Expect(svc.WorkerDialer).ToNot(BeNil()) + }) + + It("starts without an address when the operator gave none and the database route cannot give one", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + inst, err := cluster.NewRegistry(db).Get(context.Background(), "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(inst.AdvertisedAddr).To(BeEmpty()) + Expect(inst.PeerTokenHash).ToNot(BeEmpty()) + }) + + It("ends the tunnel that it holds for a node that was removed, and says whether it held one", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + Expect(svc.Disconnect("nobody")).To(BeFalse()) + + front, back := tunnelSessionPair() + _, err = svc.Tunnels.Attach(context.Background(), "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + Expect(svc.Disconnect("w1")).To(BeTrue()) + Eventually(back.CloseChan()).Should(BeClosed()) + _, _, err = cluster.NewRegistry(db).Owner(context.Background(), "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("closes the tunnels it holds before it leaves the instances table", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + front, back := tunnelSessionPair() + _, err = svc.Tunnels.Attach(context.Background(), "w2", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + + svc.Shutdown() + + Eventually(back.CloseChan()).Should(BeClosed()) + _, _, err = cluster.NewRegistry(db).Owner(context.Background(), "w2") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + It("starts a second replica on the same row without changing it", func() { first, err := initDistributed(cfg, db, nil, nil) Expect(err).ToNot(HaveOccurred()) @@ -165,6 +225,34 @@ var _ = Describe("distributed startup and the cluster carrier", func() { }) }) +// tunnelSessionPair returns the two ends of a tunnel session over a real +// websocket. +func tunnelSessionPair() (front, back *tunnel.Session) { + GinkgoHelper() + up := tunnel.NewUpgrader() + got := make(chan *tunnel.Session, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := tunnel.ServerSession(ws, tunnel.LaneInference) + if err != nil { + _ = ws.Close() + return + } + got <- sess + })) + DeferCleanup(srv.Close) + ws, _, err := tunnel.NewDialer(5*time.Second).Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil) + Expect(err).ToNot(HaveOccurred()) + back, err = tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = back.Close() }) + Eventually(got).Should(Receive(&front)) + return front, back +} + var ( natsOnce sync.Once natsURL string diff --git a/core/http/app.go b/core/http/app.go index 00362ce39d75..cf65e6b8427d 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -571,9 +571,7 @@ func API(application *application.Application) (*echo.Echo, error) { workerHTTPDial = d.WorkerHTTPDial tunnels = d.Tunnels registerOpts = append(registerOpts, localai.WithCarrierReader(d.Carriers)) - if d.Tunnels != nil { - registerOpts = append(registerOpts, localai.WithTunnelDisconnector(d.Tunnels)) - } + registerOpts = append(registerOpts, localai.WithTunnelDisconnector(d)) if d.Router != nil { remoteUnloader = d.Router.Unloader() } diff --git a/core/services/carrier/dial_conformance_test.go b/core/services/carrier/dial_conformance_test.go index 1a0fd077183a..c746cd1fb77b 100644 --- a/core/services/carrier/dial_conformance_test.go +++ b/core/services/carrier/dial_conformance_test.go @@ -199,3 +199,21 @@ func (fx *dialFixture) unroutedVia(r *replica) func(context.Context) (net.Conn, dial := r.set.Dialer("a-node-that-no-worker-holds") return func(ctx context.Context) (net.Conn, error) { return dial(ctx, "tcp", "nowhere.worker.invalid:80") } } + +var _ = Describe("Forgetting a node", func() { + It("is offered by the NATS set for its HTTP stager, and not by the one for shared storage", func() { + natsSet, err := carrier.NewNATSSet(carrier.NATSOptions{ + Client: &closingBus{FakeBus: testutil.NewFakeBus()}, Registry: noModels{}, + HTTPAddrFor: func(string) (string, error) { return "127.0.0.1:1", nil }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(natsSet.ForgetNode).ToNot(BeNil()) + natsSet.ForgetNode("n") + + s3Set, err := carrier.NewNATSSet(carrier.NATSOptions{ + Client: &closingBus{FakeBus: testutil.NewFakeBus()}, Registry: noModels{}, S3Staging: true, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(s3Set.ForgetNode).To(BeNil(), "the shared-storage stager keeps no client for each node") + }) +}) diff --git a/core/services/carrier/nats.go b/core/services/carrier/nats.go index 04415c0bcee8..768d898f671f 100644 --- a/core/services/carrier/nats.go +++ b/core/services/carrier/nats.go @@ -71,6 +71,11 @@ func NewNATSSet(o NATSOptions) (*Set, error) { Agents: nodes.NewNATSAgentControl(o.Client), Close: o.Client.Close, } + // The HTTP stager keeps one client for each node, and nothing else ever drops + // it. + if stager, ok := files.(*nodes.HTTPFileStager); ok { + set.ForgetNode = stager.ForgetNode + } if err := set.Validate(); err != nil { return nil, err } diff --git a/core/services/carrier/set.go b/core/services/carrier/set.go index f98e7c5e9df7..86639e1b39a9 100644 --- a/core/services/carrier/set.go +++ b/core/services/carrier/set.go @@ -52,6 +52,11 @@ type Set struct { Dialer nodes.WorkerNetDialerFor Agents mcpTools.AgentControl + // ForgetNode drops what the set caches for a node that left or lost its + // right to a connection, and closes the idle streams that it holds. It may + // be nil for a carrier that caches nothing per node. + ForgetNode func(nodeID string) + // Close releases the carrier's connections. It may be nil. Close func() } diff --git a/core/services/carrier/tunnel.go b/core/services/carrier/tunnel.go index d7499063c587..7fc0e04fcebb 100644 --- a/core/services/carrier/tunnel.go +++ b/core/services/carrier/tunnel.go @@ -44,25 +44,9 @@ type TunnelOptions struct { Token string } -// TunnelSet is a Set for the tunnel carrier, with the pieces that hold a cache -// per node. Forget drops them when a node leaves. -type TunnelSet struct { - *Set - control *nodes.ControlClient - files *nodes.HTTPFileStager -} - -// ForgetNode drops the clients that this set caches for a node and closes the -// idle streams they hold. Call it when a node is removed or loses its tunnel for -// good. -func (t *TunnelSet) ForgetNode(nodeID string) { - t.control.ForgetNode(nodeID) - t.files.ForgetNode(nodeID) -} - // NewTunnelSet builds the set of seam implementations that reach workers through // their tunnels. The set is complete or it is an error, as for NATS. -func NewTunnelSet(o TunnelOptions) (*TunnelSet, error) { +func NewTunnelSet(o TunnelOptions) (*Set, error) { if o.Dialer == nil { return nil, errors.New("tunnel set needs a dialer") } @@ -94,12 +78,17 @@ func NewTunnelSet(o TunnelOptions) (*TunnelSet, error) { Clients: nodes.NewDialerClientFactory(o.Token, grpcDialerFor(o.Dialer)), Dialer: httpDial, Agents: o.Agents, - Close: o.Fanout.Close, + // The control client and the stager keep one client for each node. + ForgetNode: func(nodeID string) { + control.ForgetNode(nodeID) + files.ForgetNode(nodeID) + }, + Close: o.Fanout.Close, } if err := set.Validate(); err != nil { return nil, err } - return &TunnelSet{Set: set, control: control, files: files}, nil + return set, nil } // httpDialerFor returns the dialer of the HTTP server of a worker. diff --git a/core/services/carrier/tunnel_integration_test.go b/core/services/carrier/tunnel_integration_test.go index 6041aa649a77..27b062c45dd4 100644 --- a/core/services/carrier/tunnel_integration_test.go +++ b/core/services/carrier/tunnel_integration_test.go @@ -160,7 +160,7 @@ type replica struct { tunnels *tunnel.Registry sessions *tunnel.PeerSessions pool *tunnel.PeerPool - set *carrier.TunnelSet + set *carrier.Set active *atomic.Pointer[carrier.Set] commands *carrier.Commands @@ -207,7 +207,7 @@ func startReplica(ctx context.Context, id string, db *gorm.DB, dsn string, clust Expect(err).ToNot(HaveOccurred()) r.set = set r.active = &atomic.Pointer[carrier.Set]{} - r.active.Store(set.Set) + r.active.Store(set) r.commands = carrier.NewCommands(r.active) r.files = carrier.NewFiles(r.active) r.clients = carrier.NewClients(r.active) @@ -453,6 +453,13 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { got, err := os.ReadFile(remote) Expect(err).ToNot(HaveOccurred()) Expect(got).To(Equal(payload), r.id) + + // A node that is forgotten loses its cached clients, and the next + // transfer builds new ones. + Expect(r.set.ForgetNode).ToNot(BeNil()) + r.set.ForgetNode(nodeID) + _, err = r.files.EnsureRemote(ctx, nodeID, src, "models/"+r.id+"/again.bin") + Expect(err).ToNot(HaveOccurred(), r.id) } }) From 089f9d9ea4b694c343d44c1e06f50bebae67087a Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:24:14 +0000 Subject: [PATCH 067/162] test(nodes): hold the carriers of the load verbs to two more conditions Every carrier of LoadOperationControl now also has to keep the budget of the caller running out apart from a missing route, and to read a reply that is not a reply as an error that is neither a missing route nor the answer of a worker. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../nodes/control_tunnel_conformance_test.go | 12 +++++++++ ...load_operation_control_conformance_test.go | 27 +++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/core/services/nodes/control_tunnel_conformance_test.go b/core/services/nodes/control_tunnel_conformance_test.go index 665fcdf5eddf..2f4e1c00914f 100644 --- a/core/services/nodes/control_tunnel_conformance_test.go +++ b/core/services/nodes/control_tunnel_conformance_test.go @@ -79,6 +79,8 @@ func (h *httpLoadOperationHarness) WorkerRefuses() { h.setMode("refuse") } func (h *httpLoadOperationHarness) WorkerAnswers() { h.setMode("answer") } +func (h *httpLoadOperationHarness) WorkerAnswersUnreadably() { h.setMode("unreadable") } + func (h *httpLoadOperationHarness) setMode(m string) { h.mu.Lock() defer h.mu.Unlock() @@ -148,6 +150,16 @@ func (h *httpLoadOperationHarness) serve(w http.ResponseWriter, r *http.Request) h.bodies = append(h.bodies, rec) h.mu.Unlock() + if mode == "unreadable" { + if rec.Verb == "install" { + w.Header().Set("Content-Type", workerctl.ContentTypeStream) + _, _ = io.WriteString(w, "this is not a line of the stream\n") + return + } + _, _ = io.WriteString(w, `"this is not a reply"`) + return + } + if mode == "hang" { <-r.Context().Done() return diff --git a/core/services/nodes/load_operation_control_conformance_test.go b/core/services/nodes/load_operation_control_conformance_test.go index 615681f3a7fb..de5bc75b2460 100644 --- a/core/services/nodes/load_operation_control_conformance_test.go +++ b/core/services/nodes/load_operation_control_conformance_test.go @@ -40,6 +40,9 @@ type loadOperationHarness interface { WorkerRefuses() // WorkerAnswers makes the worker answer every call with success. WorkerAnswers() + // WorkerAnswersUnreadably makes the worker answer every call with a body + // that is no reply of the verb. + WorkerAnswersUnreadably() // Sent returns the requests the carrier sent, in order. Sent() []sentRequest } @@ -89,6 +92,12 @@ func (h *natsLoadOperationHarness) WorkerRefuses() { h.mc.scriptReply(messaging.SubjectNodeModelUnload(node), workerctl.ModelUnloadReply{Success: false, Error: "process was replaced during unload"}) } +func (h *natsLoadOperationHarness) WorkerAnswersUnreadably() { + for _, s := range h.subjects() { + h.mc.scriptReply(s, "this is not a reply") + } +} + func (h *natsLoadOperationHarness) WorkerAnswers() { const node = conformanceNode h.mc.scriptReply(messaging.SubjectNodeBackendInstall(node), workerctl.BackendInstallReply{Success: true, Address: "127.0.0.1:9001", ProcessInstance: "i", ReportsOperations: true}) @@ -178,6 +187,24 @@ var _ = Describe("LoadOperationControl conformance", func() { } }) + It("does not report the budget of the caller running out as ErrNoRoute", func() { + h.TimesOut() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := h.Control().StopLoadOperation(ctx, conformanceNode, workerctl.ModelStopRequest{ProcessKey: "m#0", OperationID: "op"}) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "an impatient caller says nothing about the node: %v", err) + }) + + It("does not take an unreadable reply for a route that is missing, or for the answer of a worker", func() { + h.WorkerAnswersUnreadably() + for _, c := range calls() { + err := c.run() + Expect(err).To(HaveOccurred(), c.verb) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "%s: %v", c.verb, err) + } + }) + It("never reports a worker's own refusal as ErrNoRoute", func() { h.WorkerRefuses() for _, c := range calls() { From bec673d55b3561112627d9eeb84836e0acb7ef85 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:24:30 +0000 Subject: [PATCH 068/162] docs(distributed): add the flag for the address that peers dial A frontend publishes the address at which the other frontends dial it, for the replica that does not hold the connection of a worker to reach the one that does. The flag overrides the address that is found from the route to the database. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- docs/content/features/distributed-mode.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 44c091a0163c..b29579c05a78 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -63,6 +63,7 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These |------|---------|---------|-------------| | `--distributed` | `LOCALAI_DISTRIBUTED` | `false` | Enable distributed mode | | `--instance-id` | `LOCALAI_INSTANCE_ID` | auto UUID | Unique instance ID for this frontend | +| `--peer-address` | `LOCALAI_PEER_ADDRESS` | *(found from the route to the database)* | `host:port` at which the other frontends dial this one. A frontend publishes it so that a replica that does not hold the connection of a worker can reach the replica that does. Set it when the address found from the route to the database is not the one the other frontends should use. | | `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS server URL (e.g., `nats://localhost:4222`) | | `--registration-token` | `LOCALAI_REGISTRATION_TOKEN` | *(empty)* | Token that workers must provide to register | | `--registration-require-auth` | `LOCALAI_REGISTRATION_REQUIRE_AUTH` | `false` | Fail startup when distributed mode is enabled but the registration token is empty (node endpoints and worker file-transfer would otherwise be unauthenticated) | From 4944568c3c3c1fd5db80e639a47693fae058604b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:25:43 +0000 Subject: [PATCH 069/162] test(carrier): compare the probe on two lanes instead of timing it against a constant The spec that keeps a small call fast during a transfer ran against a fixed bound, which a loaded runner can break. It now runs the same transfer on the bulk lane and on the lane of the probe, and asks that the probe waits less behind the first. Both runs share the machine, so load moves both. A loose bound stays as a sanity check. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../carrier/tunnel_integration_test.go | 84 ++++++++++++------- 1 file changed, 52 insertions(+), 32 deletions(-) diff --git a/core/services/carrier/tunnel_integration_test.go b/core/services/carrier/tunnel_integration_test.go index 27b062c45dd4..313dd3417b2d 100644 --- a/core/services/carrier/tunnel_integration_test.go +++ b/core/services/carrier/tunnel_integration_test.go @@ -478,41 +478,61 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { src := filepath.Join(GinkgoT().TempDir(), "big.bin") Expect(os.WriteFile(src, make([]byte, 16<<20), 0o600)).To(Succeed()) - var ( - samples []time.Duration - stop = make(chan struct{}) - wg sync.WaitGroup - ) - wg.Go(func() { - for { - select { - case <-stop: - return - case <-time.After(10 * time.Millisecond): - } - begun := time.Now() - if _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"op"}}); err != nil { - return + // probeDuring probes the inference lane while transfer runs, and returns + // the 99th percentile of the probes and the time of the transfer. + probeDuring := func(transfer func() error) (time.Duration, time.Duration) { + GinkgoHelper() + var ( + samples []time.Duration + stop = make(chan struct{}) + wg sync.WaitGroup + ) + wg.Go(func() { + for { + select { + case <-stop: + return + case <-time.After(10 * time.Millisecond): + } + begun := time.Now() + if _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"op"}}); err != nil { + return + } + samples = append(samples, time.Since(begun)) } - samples = append(samples, time.Since(begun)) - } - }) - begun := time.Now() - _, err = a.files.EnsureRemote(ctx, nodeID, src, "models/big.bin") - elapsed := time.Since(begun) - close(stop) - wg.Wait() - Expect(err).ToNot(HaveOccurred()) + }) + begun := time.Now() + err := transfer() + elapsed := time.Since(begun) + close(stop) + wg.Wait() + Expect(err).ToNot(HaveOccurred()) + Expect(elapsed).To(BeNumerically(">", 500*time.Millisecond), "the transfer must last long enough to be measured") + Expect(len(samples)).To(BeNumerically(">=", 10), "probes completed during the transfer") + slices.Sort(samples) + return samples[min(len(samples)-1, len(samples)*99/100)], elapsed + } - Expect(elapsed).To(BeNumerically(">", 500*time.Millisecond), "the transfer must last long enough to be measured") - Expect(len(samples)).To(BeNumerically(">=", 10), "probes completed during the transfer") - slices.Sort(samples) - p99 := samples[min(len(samples)-1, len(samples)*99/100)] - AddReportEntry("probe during a bulk transfer", map[string]any{ - "transfer s": elapsed.Seconds(), "probes": len(samples), "p50 ms": samples[len(samples)/2].Milliseconds(), "p99 ms": p99.Milliseconds(), + // The transfer on the bulk lane, as the set does it. + bulkP99, bulkTime := probeDuring(func() error { + _, err := a.files.EnsureRemote(ctx, nodeID, src, "models/big.bin") + return err + }) + // The same transfer on the lane of the probe, which is what the set + // would do if it fell back. The comparison is between two runs on the + // same machine, so a loaded runner moves both. + shared := nodes.NewHTTPFileStager(func(id string) (string, error) { return nodes.WorkerHTTPHost(id, ""), nil }, registrationToken, a.set.Dialer) + sharedP99, sharedTime := probeDuring(func() error { + _, err := shared.EnsureRemote(ctx, nodeID, src, "models/big-shared.bin") + return err + }) + AddReportEntry("probe during a transfer", map[string]any{ + "bulk lane p99 ms": bulkP99.Milliseconds(), "bulk transfer s": bulkTime.Seconds(), + "shared lane p99 ms": sharedP99.Milliseconds(), "shared transfer s": sharedTime.Seconds(), }) - Expect(p99).To(BeNumerically("<", 100*time.Millisecond), - "a probe on the inference lane must not wait behind the bytes of a transfer on the bulk lane") + Expect(bulkP99).To(BeNumerically("<", sharedP99*8/10), + "a probe must wait less behind a transfer on the bulk lane than behind one on its own lane: bulk %v, shared %v", bulkP99, sharedP99) + Expect(bulkP99).To(BeNumerically("<", 500*time.Millisecond), "a sanity bound for the probe on a loaded machine") }) It("refuses a transfer when the worker has no bulk session, and sends nothing on the other lane", func() { From d8dc796c0170935eac83463c1211a898d765b911 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:33:28 +0000 Subject: [PATCH 070/162] test(nodes): wait for the shorter hold after a stop is acknowledged ConfirmLoadOp records the confirmation and then shortens the hold, in two updates. The spec waited for the confirmation and then read the hold at once, so under load it saw the confirmation without the shorter hold. It now waits for the hold as well. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/load_worker_ops_test.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/core/services/nodes/load_worker_ops_test.go b/core/services/nodes/load_worker_ops_test.go index 2f0a3ca16546..efc6cf1b084b 100644 --- a/core/services/nodes/load_worker_ops_test.go +++ b/core/services/nodes/load_worker_ops_test.go @@ -193,8 +193,11 @@ var _ = Describe("Load operations on the worker", func() { Expect(stops).To(HaveLen(1)) Expect(stops[0].OperationID).To(Equal(job.Generation)) Expect(stops[0].ProcessKey).To(Equal("ops-timeout#0")) - Expect(secondsUntil("ops-timeout", "stop_deadline")).To(BeNumerically("<=", loadJobFailureReport.Seconds()+3), - "an acknowledged stop frees the model in seconds, not after the full stop window") + // The confirmation and the shorter hold are two updates, so a read between + // them sees the first without the second. + Eventually(func() float64 { return secondsUntil("ops-timeout", "stop_deadline") }, 5*time.Second, 50*time.Millisecond). + Should(BeNumerically("<=", loadJobFailureReport.Seconds()+3), + "an acknowledged stop frees the model in seconds, not after the full stop window") // A request inside the report window reads the cause as a 503 answer. err := route("ops-timeout") From 43bd40bf1c6ec4757ab21ea6ceb2cd49dadd47bf Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 01:36:46 +0000 Subject: [PATCH 071/162] test(http): read the transport of the logs proxy instead of the environment Go reads the proxy variables of the environment once for the process, so a spec that sets them depends on what ran before it. The spec now looks at the transport that the logs proxy builds: it has a proxy for a worker with an address and none for a worker without one. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/localai/nodes.go | 23 ++++++--- .../localai/nodes_logs_tunnel_test.go | 50 +++++++------------ 2 files changed, 33 insertions(+), 40 deletions(-) diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index acee0624bb00..ba1f09c60f3c 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -1553,6 +1553,20 @@ func workerLogsHost(node *nodes.BackendNode) (string, bool) { return nodes.WorkerHTTPHost(node.ID, node.HTTPAddress), true } +// workerTransport returns the transport that reaches the HTTP server of a worker +// through the dialer of the node. +func workerTransport(dialFor nodes.WorkerNetDialerFor, nodeID, host string) *http.Transport { + t := httpclient.HardenedTransport() + t.DialContext = dialFor(nodeID) + // A worker without an address is reached through a tunnel, and a proxy of the + // environment cannot carry that: the transport would dial the proxy through + // the dialer of the worker. The proxy stays for every other worker. + if nodes.IsTunnelOnlyHost(host) { + t.Proxy = nil + } + return t +} + // proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer token auth. // The connection goes through dialFor(nodeID) because the advertised address // alone does not say how this frontend reaches that worker. @@ -1569,14 +1583,7 @@ func proxyHTTPToWorker(ctx context.Context, dialFor nodes.WorkerNetDialerFor, no req.Header.Set("Authorization", "Bearer "+token) } - t := httpclient.HardenedTransport() - t.DialContext = dialFor(nodeID) - // A worker without an address is reached through a tunnel, and a proxy of the - // environment cannot carry that: the transport would dial the proxy through - // the dialer of the worker. The proxy stays for every other worker. - if nodes.IsTunnelOnlyHost(httpAddress) { - t.Proxy = nil - } + t := workerTransport(dialFor, nodeID, httpAddress) client := httpclient.NewWithTimeout(15*time.Second, httpclient.WithTransport(t)) return client.Do(req) } diff --git a/core/http/endpoints/localai/nodes_logs_tunnel_test.go b/core/http/endpoints/localai/nodes_logs_tunnel_test.go index 09187f28b82e..cff1ba11204c 100644 --- a/core/http/endpoints/localai/nodes_logs_tunnel_test.go +++ b/core/http/endpoints/localai/nodes_logs_tunnel_test.go @@ -6,7 +6,6 @@ import ( "net" "net/http" "net/http/httptest" - "strings" "sync" . "github.com/onsi/ginkgo/v2" @@ -37,27 +36,26 @@ var _ = Describe("Backend logs of a worker that holds a tunnel", func() { }) Describe("the proxy of the environment", func() { - var ( - srv *httptest.Server - mu sync.Mutex - dials []string - ) + dial := func(string) func(context.Context, string, string) (net.Conn, error) { return nil } - BeforeEach(func() { - dials = nil - srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = io.WriteString(w, "logs for "+r.Host) - })) - DeferCleanup(srv.Close) - // A proxy that nothing listens on. A request that honours it fails. - GinkgoT().Setenv("HTTP_PROXY", "http://127.0.0.1:1") - GinkgoT().Setenv("http_proxy", "http://127.0.0.1:1") - GinkgoT().Setenv("NO_PROXY", "") - GinkgoT().Setenv("no_proxy", "") + It("is not used to reach a worker through its tunnel, so the dialer gets the stream", func() { + Expect(workerTransport(dial, "n1", "n1.worker.invalid").Proxy).To(BeNil()) + }) + + It("is still used for a worker that has an address, as before", func() { + Expect(workerTransport(dial, "n1", "10.0.0.1:50050").Proxy).ToNot(BeNil()) }) - dialer := func() nodes.WorkerNetDialerFor { - return func(string) func(context.Context, string, string) (net.Conn, error) { + It("sends the request of a worker without an address to the dialer even when the environment names a proxy", func() { + // The proxy of the environment is read once for the process, so this + // spec does not rely on it: it checks that the transport has none. + var ( + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "logs for "+r.Host) })) + mu sync.Mutex + dials []string + ) + DeferCleanup(srv.Close) + dialer := func(string) func(context.Context, string, string) (net.Conn, error) { return func(ctx context.Context, _, addr string) (net.Conn, error) { mu.Lock() dials = append(dials, addr) @@ -66,10 +64,7 @@ var _ = Describe("Backend logs of a worker that holds a tunnel", func() { return d.DialContext(ctx, "tcp", srv.Listener.Addr().String()) } } - } - - It("is not used to reach a worker through its tunnel, so the dialer gets the stream", func() { - resp, err := proxyHTTPToWorker(context.Background(), dialer(), "n1", "n1.worker.invalid", "/v1/backend-logs", "token") + resp, err := proxyHTTPToWorker(context.Background(), dialer, "n1", "n1.worker.invalid", "/v1/backend-logs", "token") Expect(err).ToNot(HaveOccurred()) defer func() { _ = resp.Body.Close() }() body, _ := io.ReadAll(resp.Body) @@ -79,14 +74,5 @@ var _ = Describe("Backend logs of a worker that holds a tunnel", func() { Expect(dials).To(HaveLen(1)) Expect(dials[0]).To(HavePrefix("n1.worker.invalid"), "the dialer was asked for the worker, not for a proxy") }) - - It("is still used for a worker that has an address, as before", func() { - _, err := proxyHTTPToWorker(context.Background(), dialer(), "n1", "10.0.0.1:50050", "/v1/backend-logs", "token") - Expect(err).ToNot(HaveOccurred()) - mu.Lock() - defer mu.Unlock() - Expect(dials).To(HaveLen(1)) - Expect(strings.HasPrefix(dials[0], "127.0.0.1:1")).To(BeTrue(), "with a proxy in the environment the dial goes to the proxy: %v", dials) - }) }) }) From be68af88609e6566cf59ea046775369de76adf55 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:01:33 +0000 Subject: [PATCH 072/162] feat(workerctl): name the verbs of an agent worker and let a progress line carry a subject An agent worker serves four verbs that a backend worker does not: an MCP tool call, MCP discovery, an agent run and an MCP CI run. They have paths under the control prefix like the others, and their own list so that the specs over the backend verbs do not ask a backend worker to serve them. A progress line of a streaming verb can now name a broadcast subject. A worker that holds no connection to a bus uses it to ask the frontend to publish the events of a run. A run that ends sends a RunReply as its last line. The control client hands the subject to the caller of a streaming verb. It also tells a worker that answered "no free slot" from a worker that failed: ErrWorkerBusy is not ErrNoRoute, because the route exists and the run did not start, so a caller may offer the run to another worker. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/control_client.go | 18 ++++- .../nodes/control_client_stream_test.go | 76 +++++++++++++++++++ core/services/nodes/control_tunnel.go | 4 +- core/services/workerctl/agent.go | 41 ++++++++++ core/services/workerctl/paths.go | 5 ++ core/services/workerctl/request.go | 18 +++++ core/services/workerctl/wire_test.go | 38 ++++++++++ 7 files changed, 194 insertions(+), 6 deletions(-) create mode 100644 core/services/nodes/control_client_stream_test.go create mode 100644 core/services/workerctl/agent.go diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go index 850563fc02f5..66b537665f87 100644 --- a/core/services/nodes/control_client.go +++ b/core/services/nodes/control_client.go @@ -148,12 +148,13 @@ func (c *ControlClient) Call(ctx context.Context, nodeID, verb string, req, repl // CallStreaming sends a request whose answer is a stream of lines: zero or more // progress lines and then one reply line. It calls onProgress for each progress -// line, in the order the worker sent them, and decodes the reply line into -// reply. onProgress may be nil. +// line, in the order the worker sent them, with the subject the worker named for +// it (empty for a line meant for this caller alone), and decodes the reply line +// into reply. onProgress may be nil. // // onProgress runs on the goroutine of the caller. A slow callback holds up the // body of this request and nothing else. -func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, verb string, req, reply any, onProgress func(raw json.RawMessage)) error { +func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, verb string, req, reply any, onProgress func(subject string, raw json.RawMessage)) error { resp, err := c.do(ctx, nodeID, verb, req) if err != nil { return err @@ -178,7 +179,7 @@ func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, verb string, break } if env.Progress != nil && onProgress != nil { - onProgress(env.Progress) + onProgress(env.Subject, env.Progress) } } if reply == nil { @@ -230,10 +231,19 @@ func (c *ControlClient) do(ctx context.Context, nodeID, verb string, req any) (* // is evidence about a backend. detail := readErrorBody(resp) _ = resp.Body.Close() + if workerctl.IsBusy(resp.StatusCode, detail) { + return nil, controlFailure(ctx, nodeID, fmt.Errorf("the %s verb: %w", verb, ErrWorkerBusy)) + } return nil, controlFailure(ctx, nodeID, fmt.Errorf("the %s verb answered HTTP %d: %s", verb, resp.StatusCode, detail)) } } +// ErrWorkerBusy means that the worker answered that it has no free slot for a +// run. The worker is present and the request did not start, so it says nothing +// about the work, and a caller may offer the run to another worker. It is not +// ErrNoRoute: the route exists. +var ErrWorkerBusy = errors.New("the worker has no free slot") + // maxControlErrorBodyBytes bounds how much of an error body reaches a log line. const maxControlErrorBodyBytes = 512 diff --git a/core/services/nodes/control_client_stream_test.go b/core/services/nodes/control_client_stream_test.go new file mode 100644 index 000000000000..42ce27b1f72f --- /dev/null +++ b/core/services/nodes/control_client_stream_test.go @@ -0,0 +1,76 @@ +package nodes + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +var _ = Describe("The control client and the verbs of an agent worker", func() { + const node = "11111111-2222-3333-4444-555555555555" + + var ( + worker *controlWorker + client *ControlClient + ) + + BeforeEach(func() { + worker = newControlWorker() + client = NewControlClient(worker.dialerFor(), "registration-token") + }) + + It("hands the subject of each progress line to the caller, with its bytes", func() { + worker.on(workerctl.VerbAgentExecute, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"subject":"jobs.j1.progress","progress":{"n":1}}`+"\n") + _, _ = io.WriteString(w, `{"progress":{"n":2}}`+"\n") + _, _ = io.WriteString(w, `{"reply":{"job_id":"j1","status":"completed"}}`+"\n") + }) + type line struct { + subject string + raw string + } + var seen []line + var reply workerctl.RunReply + err := client.CallStreaming(context.Background(), node, workerctl.VerbAgentExecute, map[string]string{"x": "y"}, &reply, + func(subject string, raw json.RawMessage) { seen = append(seen, line{subject, string(raw)}) }) + Expect(err).ToNot(HaveOccurred()) + Expect(seen).To(Equal([]line{{"jobs.j1.progress", `{"n":1}`}, {"", `{"n":2}`}})) + Expect(reply).To(Equal(workerctl.RunReply{JobID: "j1", Status: "completed"})) + }) + + It("tells a worker that has no free slot from a worker that failed", func() { + worker.on(workerctl.VerbMCPCIRun, func(w http.ResponseWriter, _ *http.Request) { + workerctl.WriteBusy(w) + }) + err := client.CallStreaming(context.Background(), node, workerctl.VerbMCPCIRun, struct{}{}, nil, nil) + Expect(err).To(MatchError(ErrWorkerBusy)) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "a busy worker is present and routable") + + worker.on(workerctl.VerbMCPCIRun, func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "upstream failure", http.StatusServiceUnavailable) + }) + err = client.CallStreaming(context.Background(), node, workerctl.VerbMCPCIRun, struct{}{}, nil, nil) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrWorkerBusy)).To(BeFalse(), "only the answer of the worker counts") + }) + + It("keeps the budget of the caller ahead of the answer of a busy worker", func() { + worker.on(workerctl.VerbMCPCIRun, func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(60 * time.Millisecond) + workerctl.WriteBusy(w) + }) + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + err := client.CallStreaming(ctx, node, workerctl.VerbMCPCIRun, struct{}{}, nil, nil) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + Expect(errors.Is(err, ErrWorkerBusy)).To(BeFalse()) + }) +}) diff --git a/core/services/nodes/control_tunnel.go b/core/services/nodes/control_tunnel.go index f118593a70c9..eb3a1f4b932c 100644 --- a/core/services/nodes/control_tunnel.go +++ b/core/services/nodes/control_tunnel.go @@ -70,9 +70,9 @@ func (l *httpLink) requestWithProgress(ctx context.Context, nodeID, verb, opID s } ctx, cancel := context.WithTimeout(ctx, l.bound(timeout)) defer cancel() - var sink func(json.RawMessage) + var sink func(string, json.RawMessage) if onProgress != nil && opID != "" { - sink = func(raw json.RawMessage) { + sink = func(_ string, raw json.RawMessage) { var ev workerctl.BackendInstallProgressEvent if err := json.Unmarshal(raw, &ev); err != nil { xlog.Debug("malformed backend progress event", "node", nodeID, "error", err) diff --git a/core/services/workerctl/agent.go b/core/services/workerctl/agent.go new file mode 100644 index 000000000000..e3cfff96d343 --- /dev/null +++ b/core/services/workerctl/agent.go @@ -0,0 +1,41 @@ +package workerctl + +// The verbs of an agent worker. An agent worker runs no backend processes and +// stages no files, so it serves none of AllVerbs except the stop of a backend +// (VerbBackendStop), which it answers with no body. What it serves instead is MCP +// and the runs that a frontend hands to it. +// +// They are not in AllVerbs, because a spec over that list asks the backend +// worker to serve every entry. +const ( + // VerbMCPToolExecute runs one MCP tool. Unary. + VerbMCPToolExecute = "mcp.tools.execute" + // VerbMCPDiscovery lists the MCP servers, tools, prompts and resources of a + // model. Unary. + VerbMCPDiscovery = "mcp.discovery" + // VerbAgentExecute runs one agent chat. It streams: progress lines carry the + // events of the run, and the reply line is a RunReply. + VerbAgentExecute = "agent.execute" + // VerbMCPCIRun runs one MCP CI job. It streams like VerbAgentExecute. + VerbMCPCIRun = "mcp-ci.run" +) + +// AgentVerbs returns the verbs of an agent worker once, not counting +// VerbBackendStop. +func AgentVerbs() []string { + return []string{VerbMCPToolExecute, VerbMCPDiscovery, VerbAgentExecute, VerbMCPCIRun} +} + +// RunReply is the reply line of a streaming run verb. +// +// It says what JobResultEvent says, because it is the same fact: the work +// ended, and how. On a bus the result is a message that a subscriber must +// already have. Here it is the last line of the response that the caller is +// reading, so it cannot be published to nobody. A run that names no job (an +// agent chat) has an empty JobID. +type RunReply struct { + JobID string `json:"job_id,omitempty"` + Status string `json:"status,omitempty"` // "completed", "failed" or "cancelled" + Result string `json:"result,omitempty"` + Error string `json:"error,omitempty"` +} diff --git a/core/services/workerctl/paths.go b/core/services/workerctl/paths.go index f0cc580c9821..dddcd7b1fe7a 100644 --- a/core/services/workerctl/paths.go +++ b/core/services/workerctl/paths.go @@ -60,6 +60,11 @@ func PathOf(verb string) string { // the body. A caller that reads to the end of the body without a Reply line // learned nothing about the work: the link broke. type Envelope struct { + // Subject, when set on a progress line, is the broadcast subject that the + // worker asks the receiving frontend to publish the line on. A worker that + // holds no connection to a bus uses it for the events of a run. The frontend + // publishes only the subjects that the type of the worker may ask for. + Subject string `json:"subject,omitempty"` Progress json.RawMessage `json:"progress,omitempty"` Reply json.RawMessage `json:"reply,omitempty"` } diff --git a/core/services/workerctl/request.go b/core/services/workerctl/request.go index e0391c0a15f0..dbd9e6fbba52 100644 --- a/core/services/workerctl/request.go +++ b/core/services/workerctl/request.go @@ -3,6 +3,7 @@ package workerctl import ( "io" "net/http" + "strings" "unicode/utf8" ) @@ -53,6 +54,23 @@ func WriteUnknownPath(w http.ResponseWriter, r *http.Request) { http.Error(w, "unknown worker control path "+truncate(r.URL.Path, MaxEchoedPathBytes), http.StatusNotFound) } +// BusyMessage is the body of the answer of a worker that has no free slot for a +// run. Together with the status 503 it is how a frontend tells a worker that is +// full from a worker that failed. +const BusyMessage = "the worker has no free slot for this run" + +// WriteBusy answers a run that this worker cannot take now. The worker is +// present and routable, and the request was not started, so the frontend may +// offer it to another worker. +func WriteBusy(w http.ResponseWriter) { + http.Error(w, BusyMessage, http.StatusServiceUnavailable) +} + +// IsBusy reports whether an answer is the one WriteBusy writes. +func IsBusy(status int, body string) bool { + return status == http.StatusServiceUnavailable && strings.Contains(body, BusyMessage) +} + // truncate bounds a string that a caller chose and that is about to be echoed. // It cuts on a rune boundary, because a cut inside a multi-byte rune travels as // a replacement character through every log that reads it. diff --git a/core/services/workerctl/wire_test.go b/core/services/workerctl/wire_test.go index 319525686771..6a346149a9af 100644 --- a/core/services/workerctl/wire_test.go +++ b/core/services/workerctl/wire_test.go @@ -151,3 +151,41 @@ var _ = Describe("Control payload wire format", func() { Entry("FileReleaseReply, error", workerctl.FileReleaseReply{Error: "e"}, `{"error":"e"}`), ) }) + +var _ = Describe("Agent worker verbs", func() { + It("names four verbs that the backend worker does not serve", func() { + Expect(workerctl.AgentVerbs()).To(ConsistOf( + workerctl.VerbMCPToolExecute, workerctl.VerbMCPDiscovery, + workerctl.VerbAgentExecute, workerctl.VerbMCPCIRun, + )) + for _, v := range workerctl.AgentVerbs() { + Expect(workerctl.AllVerbs()).ToNot(ContainElement(v), v) + } + }) + + It("maps each verb onto its own path below the prefix", func() { + Expect(workerctl.PathOf(workerctl.VerbMCPToolExecute)).To(Equal("/v1/control/mcp/tools/execute")) + Expect(workerctl.PathOf(workerctl.VerbMCPDiscovery)).To(Equal("/v1/control/mcp/discovery")) + Expect(workerctl.PathOf(workerctl.VerbAgentExecute)).To(Equal("/v1/control/agent/execute")) + Expect(workerctl.PathOf(workerctl.VerbMCPCIRun)).To(Equal("/v1/control/mcp-ci/run")) + }) + + It("carries the subject of a progress line only when there is one", func() { + raw, err := json.Marshal(workerctl.Envelope{Subject: "jobs.j1.progress", Progress: json.RawMessage(`{"a":1}`)}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(raw)).To(Equal(`{"subject":"jobs.j1.progress","progress":{"a":1}}`)) + + raw, err = json.Marshal(workerctl.Envelope{Progress: json.RawMessage(`{"a":1}`)}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(raw)).To(Equal(`{"progress":{"a":1}}`)) + }) + + It("pins the bytes of the terminal reply of a run", func() { + raw, err := json.Marshal(workerctl.RunReply{JobID: "j", Status: "completed", Result: "r", Error: "e"}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(raw)).To(Equal(`{"job_id":"j","status":"completed","result":"r","error":"e"}`)) + raw, err = json.Marshal(workerctl.RunReply{}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(raw)).To(Equal(`{}`)) + }) +}) From a1408e88cfbcecc0a77b0bc22859f6463adbf417 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:01:44 +0000 Subject: [PATCH 073/162] feat(nodes): publish the broadcasts that a worker asks for, within an allow list A worker that holds a tunnel has no connection to a bus. When a run on it produces events, the worker writes them as progress lines that name a subject, and the frontend that drives the run publishes them. The Rebroadcaster publishes a line only when the node type of the worker may ask for that subject: an agent worker may ask for job progress, job results and agent events, and a backend worker or an unknown type may ask for nothing. A wildcard is refused, because a filter of the table would match its own text. A refused or failed publish is logged and returns false, never an error, because the outcome of the run belongs to the worker and a publish failure says nothing about it. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/messaging/subjects.go | 1 + core/services/nodes/rebroadcast.go | 122 ++++++++++++++++++++++++ core/services/nodes/rebroadcast_test.go | 83 ++++++++++++++++ 3 files changed, 206 insertions(+) create mode 100644 core/services/nodes/rebroadcast.go create mode 100644 core/services/nodes/rebroadcast_test.go diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 7205dd897a95..65a53d99f1e5 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -103,6 +103,7 @@ const ( const ( SubjectJobResultWildcard = "jobs.*.result" SubjectJobProgressWildcard = "jobs.*.progress" + SubjectAgentEventsWildcard = "agent.*.events.*" SubjectAgentCancelWildcard = "agent.*.cancel" SubjectGalleryCancelWildcard = "gallery.*.cancel" SubjectGalleryProgressWildcard = "gallery.*.progress" diff --git a/core/services/nodes/rebroadcast.go b/core/services/nodes/rebroadcast.go new file mode 100644 index 000000000000..9a3172d8ed7e --- /dev/null +++ b/core/services/nodes/rebroadcast.go @@ -0,0 +1,122 @@ +package nodes + +import ( + "encoding/json" + "strings" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/xlog" +) + +// workerBroadcastAllow lists, per node type, the subject filters a worker may +// ask a frontend to publish on its behalf. +// +// A worker that holds a tunnel has no connection to a bus, so the events of a +// run reach the other frontends because the frontend that drives the run +// publishes them. The worker names the subject, and this table decides. +// +// An absent node type and an empty list both deny everything. Nothing may be +// added to this map without a spec that a worker of that type is refused the +// subjects that were not added. +var workerBroadcastAllow = map[string][]string{ + NodeTypeAgent: { + // The constants and not the strings: a filter written out here stops + // matching the day its builder grows a token. + messaging.SubjectJobProgressWildcard, + messaging.SubjectJobResultWildcard, + messaging.SubjectAgentEventsWildcard, + }, + // A backend worker asks for no broadcasts. Spelled as an empty list rather + // than omitted, so the reader sees the decision. + NodeTypeBackend: {}, +} + +// mayBroadcast reports whether a worker of nodeType may have subject published +// on its behalf. +func mayBroadcast(nodeType, subject string) bool { + return mayBroadcastIn(workerBroadcastAllow, nodeType, subject) +} + +// mayBroadcastIn is mayBroadcast against a caller-supplied table. +// +// It exists so a spec can hold the table itself constant while varying what is +// IN it, which is the only way to state "an empty list denies" as a property of +// this function rather than as a property of today's entries. Production has +// exactly one table and mayBroadcast is the only way to reach it. +func mayBroadcastIn(table map[string][]string, nodeType, subject string) bool { + // An empty subject is a line that asked for no broadcast at all, so there + // is nothing here to authorise. Refusing it is not defensive tidiness: it + // keeps a filter that happened to match the empty string from turning every + // ordinary private progress tick into a publish. + if subject == "" || strings.ContainsAny(subject, "*>") { + // A wildcard is a filter for a subscriber. A worker that names one asks for + // a publish on every subject it matches, and a filter of the table would + // match its own text. + return false + } + // A missing node type indexes to nil and the loop then runs zero times, + // which is the same denial an empty list produces. One code path for both + // is deliberate: "we have never heard of this worker" and "this worker is + // allowed nothing" have the same safe answer, and giving them separate + // branches invites one of the two to acquire an exception. + for _, filter := range table[nodeType] { + if messaging.SubjectMatches(filter, subject) { + return true + } + } + return false +} + +// Rebroadcaster turns an allowed progress line into a broadcast. +// +// It holds a Broadcaster and not a MessagingClient because fan-out is all it +// needs and all it may have: a worker asking for a broadcast must not be able +// to reach a request/reply or a queue group through this path. +type Rebroadcaster struct { + bus messaging.Broadcaster +} + +// NewRebroadcaster returns a Rebroadcaster publishing on bus. +func NewRebroadcaster(bus messaging.Broadcaster) *Rebroadcaster { + return &Rebroadcaster{bus: bus} +} + +// The shape of Handle's answer is asserted here, not left to its callers. +// +// A refused or failed re-broadcast must never become an error, because the only +// errors that reach a scheduler from this direction are the worker's own answer +// and an unroutable peer, and a publish failure is neither. Changing the return +// to an error therefore has to fail to COMPILE in the file that states the +// rule, rather than fail a spec somewhere a reviewer might read as a test +// needing an update. +var _ func(string, string, json.RawMessage) bool = (*Rebroadcaster)(nil).Handle + +// Handle publishes raw on subject when the worker's type allows it, and returns +// false having published nothing when it does not. It never returns an error +// that a caller could mistake for the worker's answer: a refused or failed +// re-broadcast is logged and the RPC continues, because the RPC's outcome is +// the worker's verdict about the work and a publish failure says nothing about +// it. +// +// raw is published as the same JSON value that the worker wrote. +func (r *Rebroadcaster) Handle(nodeType, subject string, raw json.RawMessage) bool { + if !mayBroadcast(nodeType, subject) { + // Warn rather than Debug: this is a worker asking for something it has + // no business on, which is the shape of a compromised or mismatched + // worker and is worth seeing without turning logging up. + xlog.Warn("refusing a worker's re-broadcast request", + "nodeType", nodeType, "subject", subject) + return false + } + if r == nil || r.bus == nil { + xlog.Debug("no broadcaster to re-broadcast a worker's progress line on", + "nodeType", nodeType, "subject", subject) + return false + } + if err := r.bus.Publish(subject, raw); err != nil { + xlog.Warn("a worker's re-broadcast could not be published", + "nodeType", nodeType, "subject", subject, "error", err) + return false + } + return true +} diff --git a/core/services/nodes/rebroadcast_test.go b/core/services/nodes/rebroadcast_test.go new file mode 100644 index 000000000000..bb2134cb715a --- /dev/null +++ b/core/services/nodes/rebroadcast_test.go @@ -0,0 +1,83 @@ +package nodes + +import ( + "encoding/json" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// failingBroadcaster refuses every publish. +type failingBroadcaster struct{ messaging.Broadcaster } + +func (failingBroadcaster) Publish(string, any) error { return errors.New("bus is down") } + +var _ = Describe("The broadcasts that a worker may ask for", func() { + DescribeTable("an agent worker", + func(subject string, allowed bool) { + Expect(mayBroadcast(NodeTypeAgent, subject)).To(Equal(allowed), subject) + }, + Entry("may report job progress", "jobs.j1.progress", true), + Entry("may report a job result", "jobs.j1.result", true), + Entry("may publish the events of an agent", "agent.helper.events.alice", true), + Entry("may not cancel a job", "jobs.j1.cancel", false), + Entry("may not cancel an agent run", "agent.helper.cancel", false), + Entry("may not publish on a control root", "nodes.n1.backend.stop", false), + Entry("may not publish on a root that carries state", "state.models.delta", false), + Entry("may not publish on a wildcard", "jobs.*.result", false), + Entry("may not publish on no subject", "", false), + ) + + It("denies a backend worker and an unknown type everything", func() { + for _, subject := range []string{"jobs.j1.progress", "jobs.j1.result", "agent.a.events.u"} { + Expect(mayBroadcast(NodeTypeBackend, subject)).To(BeFalse(), subject) + Expect(mayBroadcast("unknown", subject)).To(BeFalse(), subject) + Expect(mayBroadcast("", subject)).To(BeFalse(), subject) + } + }) + + It("denies everything to a type whose list is empty, whatever the subject", func() { + table := map[string][]string{"agent": {}} + Expect(mayBroadcastIn(table, "agent", "jobs.j1.progress")).To(BeFalse()) + }) + + Describe("the rebroadcaster", func() { + It("publishes the JSON value the worker wrote, and says so", func() { + bus := testutil.NewFakeBus() + var got []byte + _, err := bus.Subscribe("jobs.j1.progress", func(b []byte) { got = b }) + Expect(err).ToNot(HaveOccurred()) + + ok := NewRebroadcaster(bus).Handle(NodeTypeAgent, "jobs.j1.progress", json.RawMessage(`{"job_id":"j1","status":"running"}`)) + Expect(ok).To(BeTrue()) + Expect(string(got)).To(MatchJSON(`{"job_id":"j1","status":"running"}`)) + }) + + It("refuses a subject outside the list and publishes nothing", func() { + bus := testutil.NewFakeBus() + ok := NewRebroadcaster(bus).Handle(NodeTypeAgent, "jobs.j1.cancel", json.RawMessage(`{}`)) + Expect(ok).To(BeFalse()) + Expect(bus.PublishCount("jobs.j1.cancel")).To(BeZero()) + }) + + It("refuses every subject for a backend worker", func() { + bus := testutil.NewFakeBus() + Expect(NewRebroadcaster(bus).Handle(NodeTypeBackend, "jobs.j1.progress", json.RawMessage(`{}`))).To(BeFalse()) + Expect(bus.PublishCount("jobs.j1.progress")).To(BeZero()) + }) + + It("reports a bus that fails as not published, and not as an error", func() { + Expect(NewRebroadcaster(failingBroadcaster{}).Handle(NodeTypeAgent, "jobs.j1.result", json.RawMessage(`{}`))).To(BeFalse()) + }) + + It("does nothing without a bus", func() { + var r *Rebroadcaster + Expect(r.Handle(NodeTypeAgent, "jobs.j1.result", json.RawMessage(`{}`))).To(BeFalse()) + Expect(NewRebroadcaster(nil).Handle(NodeTypeAgent, "jobs.j1.result", json.RawMessage(`{}`))).To(BeFalse()) + }) + }) +}) From 94bb85c31596d7589cbde4087446194b131e0782 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:01:44 +0000 Subject: [PATCH 074/162] feat(nodes): send the MCP requests of the frontend to an agent worker over its tunnel The AgentControl of the tunnel carrier replaces the queue group of NATS with a query: the selector lists the agent nodes that may take work, asks which of them hold a tunnel, and picks one, preferring a tunnel that this replica holds because that call needs no relay. A node with a failed health probe stays in, because whether the request can reach the worker is what the connection read answers. The contract is the one of the NATS implementation. A decoded reply comes back with a nil error even when its Error field is set. A request that no worker could take wraps ErrNoRoute. Only the deadline of the caller applies, never its cancellation, so a chat client that leaves does not abort a tool that is running. A request goes to a second worker, up to three, only when the failure proves that the first did not start it: no route, a worker too old to serve the verb, or a worker with no free slot. A stream that broke or a timeout is returned as it is, because the tool may have run. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/agent_control.go | 126 ++++++++ core/services/nodes/agent_control_test.go | 370 ++++++++++++++++++++++ core/services/nodes/agent_selector.go | 117 +++++++ 3 files changed, 613 insertions(+) create mode 100644 core/services/nodes/agent_control.go create mode 100644 core/services/nodes/agent_control_test.go create mode 100644 core/services/nodes/agent_selector.go diff --git a/core/services/nodes/agent_control.go b/core/services/nodes/agent_control.go new file mode 100644 index 000000000000..3af7ac3f6795 --- /dev/null +++ b/core/services/nodes/agent_control.go @@ -0,0 +1,126 @@ +package nodes + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/mudler/LocalAI/core/config" + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// maxAgentPicks bounds how many agent workers one request is offered to. A +// caller is waiting, and a fleet whose first three workers cannot take the +// request has a problem that a fourth try does not fix. A worker whose tunnel +// broke between the connection read and the dial is ordinary, and giving up on +// the first of those would turn every worker that moves between replicas into an +// outage. +const maxAgentPicks = 3 + +// ErrNoAgentControl reports that this process was built without an agent control +// client. It says nothing about any worker. +var ErrNoAgentControl = errors.New("nodes: this deployment has no agent control client") + +// AgentControlClient is the AgentControl of the tunnel carrier. It sends the MCP +// requests of the frontend to the agent worker that the selector picks, over the +// tunnel of that worker. +// +// The contract is the one of the NATS implementation. A decoded reply comes back +// with a nil error, also when its Error field is set, because that is the answer +// of the worker. A request that no agent worker could be offered wraps +// ErrNoRoute. A timeout, a broken stream and a reply that cannot be read are +// ordinary errors. +type AgentControlClient struct { + sel *AgentSelector + cc *ControlClient +} + +// NewAgentControlClient returns the client that sends the agent requests that +// sel picks a worker for over cc. +func NewAgentControlClient(sel *AgentSelector, cc *ControlClient) *AgentControlClient { + return &AgentControlClient{sel: sel, cc: cc} +} + +// ExecuteMCPTool runs one MCP tool on an agent worker. +func (a *AgentControlClient) ExecuteMCPTool(ctx context.Context, req mcpremote.MCPToolRequest) (*mcpremote.MCPToolResponse, error) { + return agentVerb[mcpremote.MCPToolRequest, mcpremote.MCPToolResponse](ctx, a, workerctl.VerbMCPToolExecute, config.DefaultMCPToolTimeout, req) +} + +// DiscoverMCPTools lists the MCP servers and tools that the configuration of a +// model reaches. +func (a *AgentControlClient) DiscoverMCPTools(ctx context.Context, req mcpremote.MCPDiscoveryRequest) (*mcpremote.MCPDiscoveryResponse, error) { + return agentVerb[mcpremote.MCPDiscoveryRequest, mcpremote.MCPDiscoveryResponse](ctx, a, workerctl.VerbMCPDiscovery, config.DefaultMCPDiscoveryTimeout, req) +} + +// agentVerb sends one request to up to maxAgentPicks workers, one after the +// other, and stops at the first answer. +// +// What counts as no answer, and may go to the next worker, is only what proves +// that the request did not start: no route to the worker, a worker too old to +// serve the verb, a worker with no free slot. A stream that broke, a reply that +// cannot be read and a timeout are not on the list, because the tool may have run +// and a second run could do it twice. +// +// Only the deadline of ctx applies, never its cancellation. A chat client that +// disconnects must not abort a tool call that is already running on a worker. +func agentVerb[Req, Rep any](ctx context.Context, a *AgentControlClient, verb string, fallback time.Duration, req Req) (*Rep, error) { + if a == nil || a.sel == nil || a.cc == nil { + // A typed nil pointer is not a nil interface, so the receiver is checked + // as well as its parts. + return nil, fmt.Errorf("control request %s: %w", verb, ErrNoAgentControl) + } + callCtx, cancel, err := agentCallContext(ctx, fallback) + if err != nil { + return nil, err + } + defer cancel() + + tried := make(map[string]bool, maxAgentPicks) + var last error + for range maxAgentPicks { + nodeID, _, err := a.sel.pickConnectedExcluding(callCtx, tried) + if err != nil { + if last != nil && errors.Is(err, ErrNoAgentWorker) { + // Every worker that holds a tunnel has failed to take the request. + // What they said is the evidence; the end of the list is not. + return nil, last + } + return nil, err + } + tried[nodeID] = true + + var reply Rep + if err := a.cc.Call(callCtx, nodeID, verb, req, &reply); err != nil { + last = err + if callerRanOut(callCtx) != nil || !requestNeverStarted(err) { + return nil, err + } + continue + } + return &reply, nil + } + return nil, last +} + +// requestNeverStarted reports whether the failure of a call proves that the +// worker did not start the request. +func requestNeverStarted(err error) bool { + return errors.Is(err, ErrNoRoute) || errors.Is(err, errVerbNotServed) || errors.Is(err, ErrWorkerBusy) +} + +// agentCallContext keeps the deadline of ctx and drops its cancellation. Without +// a deadline the call gets the default of its verb. A budget that is spent +// already is refused here, and is the budget of the caller running out. +func agentCallContext(ctx context.Context, fallback time.Duration) (context.Context, context.CancelFunc, error) { + budget := fallback + if deadline, ok := ctx.Deadline(); ok { + budget = time.Until(deadline) + if budget <= 0 { + return nil, nil, fmt.Errorf("agent request budget spent: %w", context.DeadlineExceeded) + } + } + callCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), budget) + return callCtx, cancel, nil +} diff --git a/core/services/nodes/agent_control_test.go b/core/services/nodes/agent_control_test.go new file mode 100644 index 000000000000..1a2820835076 --- /dev/null +++ b/core/services/nodes/agent_control_test.go @@ -0,0 +1,370 @@ +package nodes + +import ( + "context" + "encoding/json" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "runtime" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// fakeConnections says which nodes are connected, and to whom. +type fakeConnections struct { + mu sync.Mutex + held map[string]bool // node id -> held by the asking replica + err error + asked [][]string +} + +func (f *fakeConnections) ConnectedAmong(_ context.Context, ids []string, _ string) ([]string, []string, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.asked = append(f.asked, append([]string(nil), ids...)) + if f.err != nil { + return nil, nil, f.err + } + var held, own []string + for _, id := range ids { + byOwner, ok := f.held[id] + if !ok { + continue + } + held = append(held, id) + if byOwner { + own = append(own, id) + } + } + return held, own, nil +} + +// agentWorkers is a set of fake agent workers, one HTTP server each, reached by a +// dialer that routes by node id. +type agentWorkers struct { + mu sync.Mutex + servers map[string]*httptest.Server + calls map[string]*atomic.Int32 +} + +func newAgentWorkers() *agentWorkers { + return &agentWorkers{servers: map[string]*httptest.Server{}, calls: map[string]*atomic.Int32{}} +} + +func (a *agentWorkers) add(id string, h http.HandlerFunc) { + a.mu.Lock() + defer a.mu.Unlock() + counter := &atomic.Int32{} + a.calls[id] = counter + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + counter.Add(1) + _, _ = io.ReadAll(r.Body) + h(w, r) + })) + DeferCleanup(srv.Close) + a.servers[id] = srv +} + +func (a *agentWorkers) callsTo(id string) int { return int(a.calls[id].Load()) } + +func (a *agentWorkers) dialerFor() WorkerNetDialerFor { + return func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + a.mu.Lock() + srv, ok := a.servers[nodeID] + a.mu.Unlock() + if !ok { + return nil, ErrNoRoute + } + var d net.Dialer + return d.DialContext(ctx, "tcp", srv.Listener.Addr().String()) + } + } +} + +func replyWith(v any) http.HandlerFunc { + return func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(v) + } +} + +var _ = Describe("The agent control of the tunnel carrier", func() { + var ( + ctx context.Context + db *gorm.DB + reg *NodeRegistry + conns *fakeConnections + workers *agentWorkers + control *AgentControlClient + ids map[string]string + ) + + register := func(name, nodeType string, autoApprove bool) string { + GinkgoHelper() + n := &BackendNode{Name: name, NodeType: nodeType} + Expect(reg.Register(ctx, n, autoApprove)).To(Succeed()) + return n.ID + } + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + ctx = context.Background() + db = testutil.SetupTestDB() + var err error + reg, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + conns = &fakeConnections{held: map[string]bool{}} + workers = newAgentWorkers() + ids = map[string]string{} + sel := NewAgentSelector(reg, conns, "replica-a") + control = NewAgentControlClient(sel, NewControlClient(workers.dialerFor(), "registration-token")) + }) + + connect := func(name string, byOwner bool, h http.HandlerFunc) { + GinkgoHelper() + ids[name] = register(name, NodeTypeAgent, true) + conns.held[ids[name]] = byOwner + workers.add(ids[name], h) + } + + toolReq := mcpremote.MCPToolRequest{ModelName: "m", ToolName: "t"} + + Describe("choosing an agent worker", func() { + It("reports no agent worker, as a missing route, when none is registered", func() { + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue()) + }) + + It("reports no agent worker when none of the registered ones holds a tunnel", func() { + register("idle", NodeTypeAgent, true) + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + }) + + It("keeps a database error apart from a missing route", func() { + register("a1", NodeTypeAgent, true) + conns.err = errors.New("database is away") + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoAgentWorker)).To(BeFalse()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("prefers a worker whose tunnel this replica holds, because that call needs no relay", func() { + remote := register("remote", NodeTypeAgent, true) + own := register("own", NodeTypeAgent, true) + conns.held[remote], conns.held[own] = false, true + for range 20 { + id, nodeType, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal(own)) + Expect(nodeType).To(Equal(NodeTypeAgent)) + } + }) + + It("falls back to a worker that another replica holds", func() { + remote := register("remote", NodeTypeAgent, true) + conns.held[remote] = false + id, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(id).To(Equal(remote)) + }) + + It("offers only agent workers that may take work", func() { + backend := register("backend", NodeTypeBackend, true) + pending := register("pending", NodeTypeAgent, false) + draining := register("draining", NodeTypeAgent, true) + Expect(reg.MarkDraining(ctx, draining)).To(Succeed()) + for _, id := range []string{backend, pending, draining} { + conns.held[id] = true + } + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + Expect(conns.asked).ToNot(BeEmpty()) + for _, asked := range conns.asked { + Expect(asked).To(BeEmpty(), "a node that may take no work is not even asked about") + } + }) + + It("does not read a failed health probe as a reason to skip a worker that holds a tunnel", func() { + id := register("flaky", NodeTypeAgent, true) + Expect(reg.MarkUnhealthy(ctx, id)).To(Succeed()) + conns.held[id] = true + picked, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(picked).To(Equal(id)) + }) + + It("leaves out the workers that the caller has already tried", func() { + a := register("a", NodeTypeAgent, true) + b := register("b", NodeTypeAgent, true) + conns.held[a], conns.held[b] = true, true + picked, _, err := NewAgentSelector(reg, conns, "replica-a").pickConnectedExcluding(ctx, map[string]bool{a: true}) + Expect(err).ToNot(HaveOccurred()) + Expect(picked).To(Equal(b)) + }) + + It("answers as a missing route when it was built with nothing", func() { + var sel *AgentSelector + _, _, err := sel.PickConnected(ctx) + Expect(err).To(MatchError(ErrNoAgentWorker)) + }) + }) + + Describe("a tool call", func() { + It("returns the decoded reply with no error, even when the worker's own answer is an error", func() { + connect("a1", true, replyWith(mcpremote.MCPToolResponse{Error: "the tool refused its arguments"})) + reply, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Error).To(Equal("the tool refused its arguments")) + Expect(workers.callsTo(ids["a1"])).To(Equal(1), "an answer is never offered to a second worker") + }) + + It("sends the verb to its path with the registration token", func() { + var path, auth atomic.Value + connect("a1", true, func(w http.ResponseWriter, r *http.Request) { + path.Store(r.URL.Path) + auth.Store(r.Header.Get("Authorization")) + replyWith(mcpremote.MCPToolResponse{Result: "ok"})(w, r) + }) + reply, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Result).To(Equal("ok")) + Expect(path.Load()).To(Equal(workerctl.PathOf(workerctl.VerbMCPToolExecute))) + Expect(auth.Load()).To(Equal("Bearer registration-token")) + }) + + It("sends discovery to the discovery path", func() { + connect("a1", true, func(w http.ResponseWriter, r *http.Request) { + Expect(r.URL.Path).To(Equal(workerctl.PathOf(workerctl.VerbMCPDiscovery))) + replyWith(mcpremote.MCPDiscoveryResponse{Servers: []mcpremote.MCPServerInfo{{Name: "s"}}})(w, r) + }) + reply, err := control.DiscoverMCPTools(ctx, mcpremote.MCPDiscoveryRequest{ModelName: "m"}) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Servers).To(HaveLen(1)) + }) + + It("wraps ErrNoRoute when there is no agent worker to offer the call to", func() { + _, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).To(MatchError(ErrNoRoute)) + _, err = control.DiscoverMCPTools(ctx, mcpremote.MCPDiscoveryRequest{}) + Expect(err).To(MatchError(ErrNoRoute)) + }) + + It("offers the call to another worker when the first one has no tunnel any more", func() { + // a1 is connected on paper and has no server: the dial finds no route. + ids["gone"] = register("gone", NodeTypeAgent, true) + conns.held[ids["gone"]] = true + connect("a2", false, replyWith(mcpremote.MCPToolResponse{Result: "from a2"})) + reply, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Result).To(Equal("from a2")) + }) + + It("offers the call to another worker when the first does not serve the verb", func() { + connect("old", true, func(w http.ResponseWriter, r *http.Request) { workerctl.WriteUnknownPath(w, r) }) + connect("new", false, replyWith(mcpremote.MCPToolResponse{Result: "from new"})) + // The worker this replica holds is tried first; the older one answers 404. + reply, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Result).To(Equal("from new")) + Expect(workers.callsTo(ids["old"])).To(Equal(1)) + }) + + It("stops after three workers, and returns what the last one said", func() { + for _, n := range []string{"o1", "o2", "o3", "o4"} { + connect(n, true, func(w http.ResponseWriter, r *http.Request) { workerctl.WriteUnknownPath(w, r) }) + } + _, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).To(HaveOccurred()) + total := 0 + for _, n := range []string{"o1", "o2", "o3", "o4"} { + total += workers.callsTo(ids[n]) + } + Expect(total).To(Equal(3)) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "a worker that is too old is present, not missing") + }) + + It("does not offer a call to a second worker when the first one failed after it started", func() { + connect("broken", true, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"result":`) // cut short + }) + connect("other", false, replyWith(mcpremote.MCPToolResponse{Result: "x"})) + _, err := control.ExecuteMCPTool(ctx, toolReq) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + Expect(workers.callsTo(ids["other"])).To(BeZero(), "the tool may have run, and a second run could do it twice") + }) + + It("reads the deadline of the caller and not its cancellation", func() { + started := make(chan struct{}) + release := make(chan struct{}) + connect("slow", true, func(w http.ResponseWriter, r *http.Request) { + close(started) + <-release + replyWith(mcpremote.MCPToolResponse{Result: "finished"})(w, r) + }) + callCtx, cancel := context.WithCancel(ctx) + type outcome struct { + reply *mcpremote.MCPToolResponse + err error + } + done := make(chan outcome, 1) + go func() { + reply, err := control.ExecuteMCPTool(callCtx, toolReq) + done <- outcome{reply, err} + }() + Eventually(started).Should(BeClosed()) + cancel() // a chat client that disconnects + Consistently(done, "200ms").ShouldNot(Receive(), "the call on the worker goes on") + close(release) + var out outcome + Eventually(done).Should(Receive(&out)) + Expect(out.err).ToNot(HaveOccurred()) + Expect(out.reply.Result).To(Equal("finished")) + }) + + It("ends the call at the deadline of the caller, and blames neither the worker nor the route", func() { + connect("slow", true, func(_ http.ResponseWriter, r *http.Request) { <-r.Context().Done() }) + callCtx, cancel := context.WithTimeout(ctx, 150*time.Millisecond) + defer cancel() + _, err := control.ExecuteMCPTool(callCtx, toolReq) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("refuses at once when the budget of the caller is already spent", func() { + connect("a1", true, replyWith(mcpremote.MCPToolResponse{})) + callCtx, cancel := context.WithDeadline(ctx, time.Now().Add(-time.Second)) + defer cancel() + _, err := control.ExecuteMCPTool(callCtx, toolReq) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + Expect(workers.callsTo(ids["a1"])).To(BeZero()) + }) + + It("refuses a client built with nothing", func() { + var nilClient *AgentControlClient + _, err := nilClient.ExecuteMCPTool(ctx, toolReq) + Expect(err).To(MatchError(ErrNoAgentControl)) + _, err = NewAgentControlClient(nil, nil).DiscoverMCPTools(ctx, mcpremote.MCPDiscoveryRequest{}) + Expect(err).To(MatchError(ErrNoAgentControl)) + }) + }) +}) diff --git a/core/services/nodes/agent_selector.go b/core/services/nodes/agent_selector.go new file mode 100644 index 000000000000..dae1c8953e4f --- /dev/null +++ b/core/services/nodes/agent_selector.go @@ -0,0 +1,117 @@ +package nodes + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +// ErrNoAgentWorker reports that no agent worker currently holds a tunnel that a +// live replica can reach. It is a fact about this moment and says nothing about +// any worker, so it wraps ErrNoRoute and nothing else: no worker is marked +// unhealthy or removed because of it. +var ErrNoAgentWorker = fmt.Errorf("no agent worker holds a tunnel to this cluster: %w", ErrNoRoute) + +// AgentConnectionReader is the part of cluster.Registry that picking needs. It +// is a port so that a spec can decide which nodes are connected without staging +// the connection rows. +type AgentConnectionReader interface { + ConnectedAmong(ctx context.Context, nodeIDs []string, owner string) (held []string, heldByOwner []string, err error) +} + +var _ AgentConnectionReader = (*cluster.Registry)(nil) + +// AgentSelector picks the agent worker that gets a request. +// +// A queue group of a broker did one thing here: it chose one subscriber out of +// several. Choosing is a query. The selector asks which agent nodes may take +// work and which of them hold a tunnel, and it prefers a worker whose tunnel +// this replica holds, because the call to it needs no relay. +// +// It answers a routing question and never an absence question. A worker that is +// not connected cannot be picked, whatever the reason. +type AgentSelector struct { + registry *NodeRegistry + conns AgentConnectionReader + // self is the id of this replica, as the connection rows record an owner. + self string +} + +// NewAgentSelector returns the selector for the agent nodes that registry lists, +// reading connections through conns, for the replica self. +func NewAgentSelector(registry *NodeRegistry, conns AgentConnectionReader, self string) *AgentSelector { + return &AgentSelector{registry: registry, conns: conns, self: self} +} + +// PickConnected returns the id and the node type of an agent node whose tunnel a +// live replica holds. The type comes back because the caller that publishes the +// events of a run needs it for every line, and the selector has read the node +// rows already. +func (s *AgentSelector) PickConnected(ctx context.Context) (nodeID, nodeType string, err error) { + return s.pickConnectedExcluding(ctx, nil) +} + +// pickConnectedExcluding is PickConnected without the nodes in tried. A retry +// that could pick the worker that has just failed would be the same call again. +func (s *AgentSelector) pickConnectedExcluding(ctx context.Context, tried map[string]bool) (string, string, error) { + if s == nil || s.registry == nil || s.conns == nil { + return "", "", fmt.Errorf("this deployment has no agent selector: %w", ErrNoAgentWorker) + } + agents, err := s.registry.selectableAgentNodes(ctx) + if err != nil { + // Plain on purpose. A database that does not answer says nothing about a + // worker, so this is neither ErrNoRoute nor ErrNoAgentWorker. + return "", "", fmt.Errorf("listing the agent workers of this deployment: %w", err) + } + typeOf := make(map[string]string, len(agents)) + ids := make([]string, 0, len(agents)) + for _, a := range agents { + if tried[a.ID] { + continue + } + typeOf[a.ID] = a.NodeType + ids = append(ids, a.ID) + } + held, heldByOwner, err := s.conns.ConnectedAmong(ctx, ids, s.self) + if err != nil { + return "", "", fmt.Errorf("reading which agent workers are connected: %w", err) + } + candidates := heldByOwner + if len(candidates) == 0 { + candidates = held + } + if len(candidates) == 0 { + return "", "", fmt.Errorf("of %d agent workers registered, none is connected: %w", len(ids), ErrNoAgentWorker) + } + // Random and not round robin: a counter on one replica says nothing about the + // load, and several replicas keep several counters that agree on nothing. + // #nosec G404 -- spreads work across equivalent workers; the choice is no secret. + picked := candidates[rand.IntN(len(candidates))] + nodeType := typeOf[picked] + if nodeType == "" { + return "", "", errors.New("the connection read named an agent worker that the node read did not: " + picked) + } + return picked, nodeType, nil +} + +// selectableAgentNodes lists the agent nodes that a request may go to. +// +// A node that waits for approval is out, and so is a node that was told to stop +// taking work. Every other status stays in, on purpose: unhealthy and offline are +// written by the health monitor on its own clock, and whether a request can reach +// the worker now is what the connection read answers. Filtering on a health +// verdict here would refuse a worker that is connected and answering because a +// probe was late. +func (r *NodeRegistry) selectableAgentNodes(ctx context.Context) ([]BackendNode, error) { + var agents []BackendNode + if err := r.db.WithContext(ctx). + Where("node_type = ? AND status NOT IN ?", NodeTypeAgent, []string{StatusPending, StatusDraining}). + Order("id"). + Find(&agents).Error; err != nil { + return nil, fmt.Errorf("listing agent nodes: %w", err) + } + return agents, nil +} From d4dc46cecbfd0b0596e0ac4dba5e748bf5d11a25 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:01:44 +0000 Subject: [PATCH 075/162] feat(agents): bind the events of an agent run to the publisher of its delivery The agent-run consumer ignored the publisher that a delivery carries and published through the one bridge of the process. That binds every run to the bus the process was built with, which holds only while the process has one bus. EventBridge.WithPublisher returns a view that publishes on the publisher of one delivery and shares the cancel registry, the store and the subscriptions of the bridge. The dispatcher builds a view for each delivery. On NATS the publisher is the NATS client, so the subjects and payloads are the ones as before. On the tunnel it writes the events to the frontend that drives the run. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/agents/dispatcher.go | 65 ++++++++------- core/services/agents/dispatcher_test.go | 50 ++++++++++++ core/services/agents/events.go | 42 +++++++--- core/services/agents/events_test.go | 103 ++++++++++++++++++++++++ 4 files changed, 221 insertions(+), 39 deletions(-) diff --git a/core/services/agents/dispatcher.go b/core/services/agents/dispatcher.go index 562be1db24c2..528ea7fffe66 100644 --- a/core/services/agents/dispatcher.go +++ b/core/services/agents/dispatcher.go @@ -246,16 +246,21 @@ func (d *NATSDispatcher) Start(ctx context.Context) error { return nil } -// runDelivery ignores events: on NATS it is the same bus the process-wide -// event bridge already publishes on. An undecodable event returns nil because -// a carrier that redelivers on error would hand it back forever. -func (d *NATSDispatcher) runDelivery(ctx context.Context, payload []byte, _ messaging.Publisher) error { +// runDelivery binds the events of the run to the publisher of the delivery. On +// NATS that publisher is the bus the bridge already publishes on, so nothing +// changes there. An undecodable event returns nil because a carrier that +// redelivers on error would hand it back forever. +func (d *NATSDispatcher) runDelivery(ctx context.Context, payload []byte, events messaging.Publisher) error { var evt AgentChatEvent if err := json.Unmarshal(payload, &evt); err != nil { xlog.Error("Failed to unmarshal agent chat event", "error", err) return nil } - d.handleJob(ctx, evt) + var bridge *EventBridge + if d.eventBridge != nil { + bridge = d.eventBridge.WithPublisher(events) + } + d.handleJob(ctx, evt, bridge) return nil } @@ -269,7 +274,7 @@ func (d *NATSDispatcher) Stop() error { return nil } -func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent) { +func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent, bridge *EventBridge) { xlog.Info("Processing agent chat job", "agent", evt.AgentName, "user", evt.UserID) // Prefer config from the enriched payload (no DB needed). @@ -280,16 +285,16 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent) { cfg, err = d.configs.GetAgentConfig(evt.UserID, evt.AgentName) if err != nil { xlog.Error("Failed to load agent config", "agent", evt.AgentName, "error", err) - if d.eventBridge != nil { - d.eventBridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") + if bridge != nil { + bridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") } return } } if cfg == nil { xlog.Error("No agent config available", "agent", evt.AgentName) - if d.eventBridge != nil { - d.eventBridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") + if bridge != nil { + bridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") } return } @@ -298,12 +303,12 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent) { defer cancel() // Register cancellation - if d.eventBridge != nil { - d.eventBridge.RegisterCancel(evt.MessageID, cancel) - defer d.eventBridge.DeregisterCancel(evt.MessageID) + if bridge != nil { + bridge.RegisterCancel(evt.MessageID, cancel) + defer bridge.DeregisterCancel(evt.MessageID) } - cb := d.buildNATSCallbacks(evt) + cb := d.buildNATSCallbacks(evt, bridge) // Build execution options: skills come from the enriched NATS payload // (workers have no database access). @@ -327,9 +332,9 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent) { if execErr != nil { xlog.Error("Distributed agent execution failed", "agent", evt.AgentName, "error", execErr) - if d.eventBridge != nil { - d.eventBridge.PublishStatus(evt.AgentName, evt.UserID, "error") - d.eventBridge.PublishMessage(evt.AgentName, evt.UserID, RoleAgent, + if bridge != nil { + bridge.PublishStatus(evt.AgentName, evt.UserID, "error") + bridge.PublishMessage(evt.AgentName, evt.UserID, RoleAgent, fmt.Sprintf("Agent execution failed: %v", execErr), evt.MessageID+"-error") } return @@ -347,7 +352,7 @@ func (p *staticSkillProvider) ListSkills() ([]SkillInfo, error) { return p.skills, nil } -func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks { +func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent, bridge *EventBridge) Callbacks { // Observable tracking: build LocalAGI-compatible observable records // from cogito callbacks so the UI can render them properly. // @@ -381,7 +386,7 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks { return Callbacks{ OnStream: func(ev cogito.StreamEvent) { - if d.eventBridge == nil { + if bridge == nil { return } data := map[string]any{"timestamp": time.Now().Format(time.RFC3339)} @@ -423,7 +428,7 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks { default: return } - d.eventBridge.PublishStreamEvent(evt.AgentName, evt.UserID, data) + bridge.PublishStreamEvent(evt.AgentName, evt.UserID, data) }, OnReasoning: func(text string) { // Reasoning is buffered via OnStream @@ -433,8 +438,8 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks { }, OnToolResult: func(name, result string) { // Emit tool_result stream event for real-time UI display - if d.eventBridge != nil { - d.eventBridge.PublishStreamEvent(evt.AgentName, evt.UserID, map[string]any{ + if bridge != nil { + bridge.PublishStreamEvent(evt.AgentName, evt.UserID, map[string]any{ "type": "tool_result", "tool_name": name, "tool_result": result, @@ -450,23 +455,23 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks { obs.Completion = &coreTypes.Completion{ ActionResult: result, } - if d.eventBridge != nil { - d.eventBridge.PersistObservable(evt.AgentName, evt.UserID, "tool_result", obs) + if bridge != nil { + bridge.PersistObservable(evt.AgentName, evt.UserID, "tool_result", obs) } } }, OnStatus: func(status string) { - if d.eventBridge != nil { - d.eventBridge.PublishStatus(evt.AgentName, evt.UserID, status) + if bridge != nil { + bridge.PublishStatus(evt.AgentName, evt.UserID, status) } }, OnMessage: func(sender, content, msgID string) { - if d.eventBridge != nil { - d.eventBridge.PublishMessage(evt.AgentName, evt.UserID, sender, content, msgID) + if bridge != nil { + bridge.PublishMessage(evt.AgentName, evt.UserID, sender, content, msgID) } // On agent response, persist the root observable with completion - if sender == RoleAgent && d.eventBridge != nil { + if sender == RoleAgent && bridge != nil { rootObs.Completion = &coreTypes.Completion{ ActionResult: content, } @@ -480,7 +485,7 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent) Callbacks { }, } } - d.eventBridge.PersistObservable(evt.AgentName, evt.UserID, "chat", rootObs) + bridge.PersistObservable(evt.AgentName, evt.UserID, "chat", rootObs) } }, } diff --git a/core/services/agents/dispatcher_test.go b/core/services/agents/dispatcher_test.go index 4e4a648e00e3..100edbf4afbf 100644 --- a/core/services/agents/dispatcher_test.go +++ b/core/services/agents/dispatcher_test.go @@ -157,3 +157,53 @@ var _ = Describe("NATSDispatcher.Start", func() { Entry("bounded", 4), ) }) + +// handlerConsumer is a WorkConsumer that keeps the handler it is given, so a +// spec can make a delivery with a publisher of its choice. +type handlerConsumer struct{ handler messaging.WorkHandler } + +func (c *handlerConsumer) Consume(_ context.Context, _ messaging.WorkKind, _ int, h messaging.WorkHandler) (messaging.Subscription, error) { + c.handler = h + return noopSubscription{}, nil +} + +var _ = Describe("NATSDispatcher delivery", func() { + var ( + own *recordingPublisher + consumer *handlerConsumer + d *NATSDispatcher + ) + + BeforeEach(func() { + own = &recordingPublisher{} + consumer = &handlerConsumer{} + bridge := NewEventBridge(&broadcasterOver{Publisher: own}, nil, "worker-1") + d = NewNATSDispatcher(consumer, bridge, nil, "http://127.0.0.1:1", "", 0) + Expect(d.Start(GinkgoT().Context())).To(Succeed()) + }) + + payload := func() []byte { + raw, err := json.Marshal(AgentChatEvent{AgentName: "a1", UserID: "u1", Message: "hi", MessageID: "m1"}) + Expect(err).ToNot(HaveOccurred()) + return raw + } + + It("binds the events of a run to the publisher of the delivery, and to nothing else", func() { + events := &recordingPublisher{} + Expect(consumer.handler(GinkgoT().Context(), payload(), events)).To(Succeed()) + + subjects, payloads := events.seen() + Expect(subjects).To(Equal([]string{messaging.SubjectAgentEvents("a1", "u1")})) + evt := payloads[0].(AgentEvent) + Expect(evt.EventType).To(Equal("json_message_status")) + Expect(evt.Metadata).To(ContainSubstring("error: agent config not found")) + ownSubjects, _ := own.seen() + Expect(ownSubjects).To(BeEmpty()) + }) + + It("keeps a delivery that has no publisher on the bridge", func() { + Expect(consumer.handler(GinkgoT().Context(), payload(), nil)).To(Succeed()) + subjects, _ := own.seen() + Expect(subjects).To(HaveLen(1)) + }) +}) diff --git a/core/services/agents/events.go b/core/services/agents/events.go index 57792ef137f8..e5609a8629a8 100644 --- a/core/services/agents/events.go +++ b/core/services/agents/events.go @@ -40,12 +40,16 @@ type AgentCancelEvent struct { // EventBridge bridges agent events between NATS and SSE connections. // It enables cross-instance SSE: user connects to Frontend 1, agent runs on Frontend 2. type EventBridge struct { - nats messaging.Broadcaster + nats messaging.Broadcaster + // pub is where events are published. It is nats unless a view bound it to + // the publisher of one delivery (see WithPublisher). + pub messaging.Publisher store *AgentStore instanceID string - // Cancel registry for running agent executions - cancelRegistry messaging.CancelRegistry + // Cancel registry for running agent executions. A view shares it with the + // bridge it came from, so a cancel finds the run whichever view registered it. + cancelRegistry *messaging.CancelRegistry // Background NATS subscriptions owned by this bridge obsPersisterSub messaging.Subscription @@ -54,12 +58,32 @@ type EventBridge struct { // NewEventBridge creates a new EventBridge. func NewEventBridge(nc messaging.Broadcaster, store *AgentStore, instanceID string) *EventBridge { return &EventBridge{ - nats: nc, - store: store, - instanceID: instanceID, + nats: nc, + pub: nc, + store: store, + instanceID: instanceID, + cancelRegistry: &messaging.CancelRegistry{}, } } +// WithPublisher returns a view of the bridge that publishes its events on pub. +// The view shares the cancel registry, the store and the subscriptions of the +// bridge. A nil pub gives a view that publishes where the bridge does. +// +// A run is bound to the publisher of the delivery that started it, and not to +// the process. On NATS that publisher is the NATS client itself, and the +// subjects and payloads are the ones the bridge publishes. On the tunnel it +// writes the events as lines of the response to the frontend that drives the +// run. Binding per delivery is what lets a run keep its carrier when the +// carrier of the cluster changes while it runs. +func (b *EventBridge) WithPublisher(pub messaging.Publisher) *EventBridge { + view := *b + if pub != nil { + view.pub = pub + } + return &view +} + // PublishEvent publishes an agent event to NATS for SSE bridging. // // Timestamp is emitted in Unix milliseconds to match the local dispatcher's @@ -70,7 +94,7 @@ func NewEventBridge(nc messaging.Broadcaster, store *AgentStore, instanceID stri func (b *EventBridge) PublishEvent(agentName, userID string, evt AgentEvent) error { evt.Timestamp = time.Now().UnixMilli() subject := messaging.SubjectAgentEvents(agentName, userID) - return b.nats.Publish(subject, evt) + return b.pub.Publish(subject, evt) } // PersistObservable publishes an observable_update SSE event for real-time UI @@ -163,7 +187,7 @@ func (b *EventBridge) CancelExecution(agentName, userID, messageID string) error } // Also publish via NATS for other instances - return b.nats.Publish(messaging.SubjectAgentCancel(agentName), AgentCancelEvent{ + return b.pub.Publish(messaging.SubjectAgentCancel(agentName), AgentCancelEvent{ AgentName: agentName, UserID: userID, MessageID: messageID, @@ -201,7 +225,7 @@ func (b *EventBridge) StartObservablePersister() error { return fmt.Errorf("no store available for observable persistence") } // Subscribe to all agent events using wildcard: agent.*.events.* - sub, err := messaging.SubscribeJSON(b.nats, "agent.*.events.*", func(evt AgentEvent) { + sub, err := messaging.SubscribeJSON(b.nats, messaging.SubjectAgentEventsWildcard, func(evt AgentEvent) { if evt.EventType != "observable_update" { return } diff --git a/core/services/agents/events_test.go b/core/services/agents/events_test.go index f103f7982e57..477c53c33ec4 100644 --- a/core/services/agents/events_test.go +++ b/core/services/agents/events_test.go @@ -1,6 +1,7 @@ package agents import ( + "sync" "time" "github.com/mudler/LocalAI/core/services/messaging" @@ -76,3 +77,105 @@ var _ = Describe("EventBridge", func() { }) }) }) + +// recordingPublisher keeps what is published to it. +type recordingPublisher struct { + mu sync.Mutex + subjects []string + payloads []any +} + +func (r *recordingPublisher) Publish(subject string, data any) error { + r.mu.Lock() + defer r.mu.Unlock() + r.subjects = append(r.subjects, subject) + r.payloads = append(r.payloads, data) + return nil +} + +func (r *recordingPublisher) seen() ([]string, []any) { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.subjects...), append([]any(nil), r.payloads...) +} + +var _ = Describe("EventBridge.WithPublisher", func() { + It("publishes every kind of event on the publisher it is given, with the subject and payload of the bridge", func() { + own := &recordingPublisher{} + bound := &recordingPublisher{} + bridge := NewEventBridge(&broadcasterOver{Publisher: own}, nil, "worker-1") + view := bridge.WithPublisher(bound) + + Expect(view.PublishMessage("a1", "u1", RoleAgent, "hi", "m1")).To(Succeed()) + Expect(view.PublishStatus("a1", "u1", "processing")).To(Succeed()) + Expect(view.PublishStreamEvent("a1", "u1", map[string]any{"type": "content"})).To(Succeed()) + view.PersistObservable("a1", "u1", "chat", map[string]string{"k": "v"}) + + subjects, payloads := bound.seen() + Expect(subjects).To(HaveLen(4)) + for _, s := range subjects { + Expect(s).To(Equal(messaging.SubjectAgentEvents("a1", "u1"))) + } + types := []string{} + for _, p := range payloads { + evt, ok := p.(AgentEvent) + Expect(ok).To(BeTrue()) + Expect(evt.AgentName).To(Equal("a1")) + types = append(types, evt.EventType) + } + Expect(types).To(Equal([]string{"json_message", "json_message_status", "stream_event", "observable_update"})) + ownSubjects, _ := own.seen() + Expect(ownSubjects).To(BeEmpty(), "nothing leaks to the bus the bridge was built with") + }) + + It("publishes the same subjects and payloads as the bridge it was made from", func() { + viaBridge := &recordingPublisher{} + viaView := &recordingPublisher{} + bridge := NewEventBridge(&broadcasterOver{Publisher: viaBridge}, nil, "worker-1") + view := NewEventBridge(&broadcasterOver{Publisher: &recordingPublisher{}}, nil, "worker-1").WithPublisher(viaView) + + Expect(bridge.PublishMessage("a1", "u1", RoleUser, "hello", "m1")).To(Succeed()) + Expect(view.PublishMessage("a1", "u1", RoleUser, "hello", "m1")).To(Succeed()) + + s1, p1 := viaBridge.seen() + s2, p2 := viaView.seen() + Expect(s2).To(Equal(s1)) + e1, e2 := p1[0].(AgentEvent), p2[0].(AgentEvent) + e1.Timestamp, e2.Timestamp = 0, 0 + Expect(e2).To(Equal(e1)) + }) + + It("shares the cancel registry with the bridge", func() { + bridge := NewEventBridge(&broadcasterOver{Publisher: &recordingPublisher{}}, nil, "worker-1") + view := bridge.WithPublisher(&recordingPublisher{}) + + cancelled := false + view.RegisterCancel("m1", func() { cancelled = true }) + Expect(bridge.cancelRegistry.Cancel("m1")).To(BeTrue()) + Expect(cancelled).To(BeTrue()) + + view.RegisterCancel("m2", func() { cancelled = false }) + view.DeregisterCancel("m2") + Expect(bridge.cancelRegistry.Cancel("m2")).To(BeFalse()) + }) + + It("keeps the publisher of the bridge when it is given none", func() { + own := &recordingPublisher{} + bridge := NewEventBridge(&broadcasterOver{Publisher: own}, nil, "worker-1") + Expect(bridge.WithPublisher(nil).PublishStatus("a1", "u1", "x")).To(Succeed()) + subjects, _ := own.seen() + Expect(subjects).To(HaveLen(1)) + }) +}) + +// broadcasterOver is a Broadcaster whose Publish goes to a recording publisher +// and which subscribes to nothing. +type broadcasterOver struct{ Publisher messaging.Publisher } + +func (b *broadcasterOver) Publish(subject string, data any) error { + return b.Publisher.Publish(subject, data) +} + +func (b *broadcasterOver) Subscribe(string, func([]byte)) (messaging.Subscription, error) { + return &fakeSub{}, nil +} From 1a77c92744eaaf686874c416bea7aff15eefc112 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:04:08 +0000 Subject: [PATCH 076/162] feat(agentworker): serve the control plane of an agent worker over its tunnel An agent worker used to subscribe to NATS for everything it did. Its MCP requests, its agent runs and its MCP CI jobs now arrive as requests on a small HTTP server that binds loopback only and is reached through the tunnel of the worker. No flag moves the bind, and the tunnel offers the http tag alone, so the frontend has no way to name a port of the host. Work is the WorkConsumer of this carrier. A handler that was written for a queue runs on it unchanged: the request for a run is the delivery, the events that the handler publishes are the lines of the response, and the result it publishes becomes the reply line. A frontend picks the worker, so a worker that is full answers "busy" instead of holding the run, and the frontend offers it to another worker. A response that ends without a reply line is a link that broke and says nothing about the run. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../agentworker/agentworker_suite_test.go | 13 + core/services/agentworker/doc.go | 16 + core/services/agentworker/handlers.go | 148 ++++++ core/services/agentworker/server_test.go | 420 ++++++++++++++++++ core/services/agentworker/start.go | 132 ++++++ core/services/agentworker/start_test.go | 95 ++++ core/services/agentworker/work.go | 263 +++++++++++ core/services/nodes/file_transfer_server.go | 13 + core/services/worker/tunnel.go | 10 + core/services/worker/tunnel_test.go | 18 + 10 files changed, 1128 insertions(+) create mode 100644 core/services/agentworker/agentworker_suite_test.go create mode 100644 core/services/agentworker/doc.go create mode 100644 core/services/agentworker/handlers.go create mode 100644 core/services/agentworker/server_test.go create mode 100644 core/services/agentworker/start.go create mode 100644 core/services/agentworker/start_test.go create mode 100644 core/services/agentworker/work.go diff --git a/core/services/agentworker/agentworker_suite_test.go b/core/services/agentworker/agentworker_suite_test.go new file mode 100644 index 000000000000..d7a63fe4c702 --- /dev/null +++ b/core/services/agentworker/agentworker_suite_test.go @@ -0,0 +1,13 @@ +package agentworker_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAgentWorker(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Agent worker test suite") +} diff --git a/core/services/agentworker/doc.go b/core/services/agentworker/doc.go new file mode 100644 index 000000000000..32b1bb3ce874 --- /dev/null +++ b/core/services/agentworker/doc.go @@ -0,0 +1,16 @@ +// Package agentworker serves the control plane of an agent worker over the +// tunnel it holds to a frontend. +// +// An agent worker runs no backend processes and stages no files. It creates the +// MCP sessions that a frontend cannot (a stdio server under docker), runs tools +// against them, answers discovery, and executes the agent runs and MCP CI jobs +// that the frontend hands to it. On NATS each of these arrived on a subject and +// a queue group picked the worker. On the tunnel the frontend picks the worker +// and sends the request over the stream it opens, so this package answers +// requests and chooses nothing. +// +// A run is a streaming request. The events of the run are the lines of the +// response, and the last line is the reply. The reply is the result of the run, +// so it cannot be published to nobody, and a response that ends without a reply +// line means the link broke: it says nothing about the work. +package agentworker diff --git a/core/services/agentworker/handlers.go b/core/services/agentworker/handlers.go new file mode 100644 index 000000000000..17182dd9184f --- /dev/null +++ b/core/services/agentworker/handlers.go @@ -0,0 +1,148 @@ +package agentworker + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// UnaryHandler answers a verb with one reply and no progress. The bytes it +// returns are the whole body of the response; nil is a 204. +// +// An error is the worker failing to serve the verb, never the bad news of the +// verb. A body on a 200 is the answer of the worker, and everything else is a +// failure that the frontend must not read as evidence about anything. An MCP +// tool that ran and failed is an answer, and travels as bytes with its Error +// field set. A request that cannot be decoded is ErrUndecodable. +type UnaryHandler func(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) + +// ErrUndecodable marks a request body that is not a request of the verb. It is +// a 400: a body that could not be read is not an answer about anything. +var ErrUndecodable = errors.New("the request body cannot be decoded") + +// Unary adapts a typed handler to a UnaryHandler. The typed handler returns no +// error, because every failure of an MCP request is an answer that the reply +// carries in its Error field, so the requester never waits out its budget on +// silence. +func Unary[Req, Rep any](h func(ctx context.Context, req Req) Rep) UnaryHandler { + return func(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) { + var req Req + if err := json.Unmarshal(raw, &req); err != nil { + return nil, fmt.Errorf("%w: %v", ErrUndecodable, err) + } + out, err := json.Marshal(h(ctx, req)) + if err != nil { + return nil, fmt.Errorf("encoding the reply: %w", err) + } + return out, nil + } +} + +// Config is the part of the control plane that answers with one reply. +// +// A nil field means that this worker does not serve the verb, and nothing is +// mounted for it, so the catch-all under the control prefix answers 404. That is +// the answer of an older build, which the frontend reads as a worker that does +// not serve the verb, and not as a worker that is gone. +type Config struct { + // MCPTool answers workerctl.VerbMCPToolExecute. + MCPTool UnaryHandler + // MCPDiscovery answers workerctl.VerbMCPDiscovery. + MCPDiscovery UnaryHandler + // BackendStop drops the MCP sessions cached for a backend that goes away. It + // answers workerctl.VerbBackendStop on the same path that a backend worker + // serves with another implementation, so the frontend sends one request to + // either kind of worker. + BackendStop func(ctx context.Context, req workerctl.BackendStopRequest) error +} + +// Handler returns the control plane of an agent worker: the verbs of cfg, the +// runs that work consumes, and the catch-all for every other path under the +// prefix. It does no authentication. Start puts one check in front of it. +func Handler(cfg Config, work *Work) http.Handler { + mux := http.NewServeMux() + if cfg.MCPTool != nil { + mux.HandleFunc(workerctl.PathOf(workerctl.VerbMCPToolExecute), serveUnary(workerctl.VerbMCPToolExecute, cfg.MCPTool)) + } + if cfg.MCPDiscovery != nil { + mux.HandleFunc(workerctl.PathOf(workerctl.VerbMCPDiscovery), serveUnary(workerctl.VerbMCPDiscovery, cfg.MCPDiscovery)) + } + if cfg.BackendStop != nil { + mux.HandleFunc(workerctl.PathOf(workerctl.VerbBackendStop), serveBackendStop(cfg.BackendStop)) + } + if work != nil { + work.register(mux) + } + // Mounted last and always: a path under the prefix that no verb claims is a + // frontend newer than this worker. + mux.HandleFunc(workerctl.Prefix, workerctl.WriteUnknownPath) + return mux +} + +// failToServe answers a verb that this worker could not serve. 500 and not 404, +// which the frontend reads as a verb the worker does not implement, and not a +// 2xx, which it reads as an answer. +func failToServe(w http.ResponseWriter, verb string, err error) { + http.Error(w, "the "+verb+" verb could not be served: "+err.Error(), http.StatusInternalServerError) +} + +// refuseRequest answers a request that this worker could not read. The fault is +// the caller's, and the status is not 2xx for the reason above. +func refuseRequest(w http.ResponseWriter, verb string, err error) { + http.Error(w, "the "+verb+" verb could not be served: "+err.Error(), http.StatusBadRequest) +} + +func serveUnary(verb string, h UnaryHandler) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + body, ok := workerctl.ReadRequestBody(w, r) + if !ok { + return + } + reply, err := h(r.Context(), json.RawMessage(body)) + switch { + case errors.Is(err, ErrUndecodable): + refuseRequest(w, verb, err) + return + case err != nil: + failToServe(w, verb, err) + return + case len(reply) == 0: + w.WriteHeader(http.StatusNoContent) + return + } + w.Header().Set("Content-Type", "application/json") + if _, err := w.Write(reply); err != nil { + xlog.Debug("agent worker reply could not be written", "verb", verb, "error", err) + } + } +} + +// serveBackendStop answers 204 with no body, which is what a backend worker +// answers on the same path. +func serveBackendStop(h func(ctx context.Context, req workerctl.BackendStopRequest) error) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + body, ok := workerctl.ReadRequestBody(w, r) + if !ok { + return + } + var req workerctl.BackendStopRequest + if err := json.Unmarshal(body, &req); err != nil { + refuseRequest(w, workerctl.VerbBackendStop, err) + return + } + // The context of the caller is dropped, as a backend worker drops it: the + // stop has no answer that anyone waits for, and a cleanup abandoned half + // way would leave sessions cached for a backend that is gone. + if err := h(context.WithoutCancel(r.Context()), req); err != nil { + failToServe(w, workerctl.VerbBackendStop, err) + return + } + w.WriteHeader(http.StatusNoContent) + } +} diff --git a/core/services/agentworker/server_test.go b/core/services/agentworker/server_test.go new file mode 100644 index 000000000000..4ac9f708914a --- /dev/null +++ b/core/services/agentworker/server_test.go @@ -0,0 +1,420 @@ +package agentworker_test + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/agentworker" + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +func post(srv *httptest.Server, verb string, body string) *http.Response { + GinkgoHelper() + resp, err := http.Post(srv.URL+workerctl.PathOf(verb), "application/json", strings.NewReader(body)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + return resp +} + +func readBody(resp *http.Response) string { + GinkgoHelper() + b, err := io.ReadAll(resp.Body) + Expect(err).ToNot(HaveOccurred()) + return string(b) +} + +var _ = Describe("The control plane of an agent worker", func() { + var work *agentworker.Work + + serve := func(cfg agentworker.Config) *httptest.Server { + GinkgoHelper() + srv := httptest.NewServer(agentworker.Handler(cfg, work)) + DeferCleanup(srv.Close) + return srv + } + + BeforeEach(func() { work = agentworker.NewWork() }) + + Describe("the unary verbs", func() { + It("answers a tool call with the reply of the handler on a 200", func() { + var got mcpremote.MCPToolRequest + srv := serve(agentworker.Config{MCPTool: agentworker.Unary(func(_ context.Context, req mcpremote.MCPToolRequest) mcpremote.MCPToolResponse { + got = req + return mcpremote.MCPToolResponse{Result: "done"} + })}) + resp := post(srv, workerctl.VerbMCPToolExecute, `{"model_name":"m","tool_name":"t"}`) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(readBody(resp)).To(MatchJSON(`{"result":"done"}`)) + Expect(got.ModelName).To(Equal("m")) + Expect(got.ToolName).To(Equal("t")) + }) + + It("answers a discovery with the reply of the handler", func() { + srv := serve(agentworker.Config{MCPDiscovery: agentworker.Unary(func(_ context.Context, req mcpremote.MCPDiscoveryRequest) mcpremote.MCPDiscoveryResponse { + return mcpremote.MCPDiscoveryResponse{Servers: []mcpremote.MCPServerInfo{{Name: req.ModelName}}} + })}) + resp := post(srv, workerctl.VerbMCPDiscovery, `{"model_name":"m"}`) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(readBody(resp)).To(ContainSubstring(`"name":"m"`)) + }) + + It("keeps the error of a tool in the reply, on a 200, because it is the answer of the worker", func() { + srv := serve(agentworker.Config{MCPTool: agentworker.Unary(func(context.Context, mcpremote.MCPToolRequest) mcpremote.MCPToolResponse { + return mcpremote.MCPToolResponse{Error: "the tool failed"} + })}) + resp := post(srv, workerctl.VerbMCPToolExecute, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(readBody(resp)).To(MatchJSON(`{"error":"the tool failed"}`)) + }) + + It("refuses a body that does not decode with 400, and never with a reply", func() { + called := false + srv := serve(agentworker.Config{MCPTool: agentworker.Unary(func(context.Context, mcpremote.MCPToolRequest) mcpremote.MCPToolResponse { + called = true + return mcpremote.MCPToolResponse{} + })}) + resp := post(srv, workerctl.VerbMCPToolExecute, `not json`) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + Expect(called).To(BeFalse()) + }) + + It("refuses anything but a POST", func() { + srv := serve(agentworker.Config{MCPTool: agentworker.Unary(func(context.Context, mcpremote.MCPToolRequest) mcpremote.MCPToolResponse { + return mcpremote.MCPToolResponse{} + })}) + resp, err := http.Get(srv.URL + workerctl.PathOf(workerctl.VerbMCPToolExecute)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusMethodNotAllowed)) + }) + + It("answers 404 with a body that says so for a verb it does not serve", func() { + srv := serve(agentworker.Config{}) + for _, verb := range []string{workerctl.VerbMCPToolExecute, workerctl.VerbMCPDiscovery, workerctl.VerbBackendInstall, workerctl.VerbModelStop} { + resp := post(srv, verb, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusNotFound), verb) + Expect(readBody(resp)).To(ContainSubstring("unknown worker control path"), verb) + } + }) + }) + + Describe("the stop of a backend", func() { + It("answers 204 with no body, as a backend worker does, and drops the sessions of the backend", func() { + var stopped []string + srv := serve(agentworker.Config{BackendStop: func(_ context.Context, req workerctl.BackendStopRequest) error { + stopped = append(stopped, req.Backend) + return nil + }}) + resp := post(srv, workerctl.VerbBackendStop, `{"backend":"llama"}`) + Expect(resp.StatusCode).To(Equal(http.StatusNoContent)) + Expect(readBody(resp)).To(BeEmpty()) + Expect(stopped).To(Equal([]string{"llama"})) + }) + + It("finishes the cleanup when the caller leaves", func() { + ctxErr := make(chan error, 1) + srv := serve(agentworker.Config{BackendStop: func(ctx context.Context, _ workerctl.BackendStopRequest) error { + time.Sleep(50 * time.Millisecond) + ctxErr <- ctx.Err() + return nil + }}) + ctx, cancel := context.WithCancel(context.Background()) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL+workerctl.PathOf(workerctl.VerbBackendStop), strings.NewReader(`{"backend":"b"}`)) + go func() { time.Sleep(10 * time.Millisecond); cancel() }() + _, _ = http.DefaultClient.Do(req) + Eventually(ctxErr).Should(Receive(BeNil())) + }) + + It("refuses a body that does not decode with 400", func() { + srv := serve(agentworker.Config{BackendStop: func(context.Context, workerctl.BackendStopRequest) error { return nil }}) + Expect(post(srv, workerctl.VerbBackendStop, `{`).StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("answers 500 when the cleanup fails, and not a reply", func() { + srv := serve(agentworker.Config{BackendStop: func(context.Context, workerctl.BackendStopRequest) error { return errors.New("boom") }}) + resp := post(srv, workerctl.VerbBackendStop, `{"backend":"b"}`) + Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError)) + }) + }) +}) + +// envelopes reads the lines of a streaming response. +func envelopes(resp *http.Response) []workerctl.Envelope { + GinkgoHelper() + var out []workerctl.Envelope + sc := bufio.NewScanner(resp.Body) + sc.Buffer(make([]byte, 1<<20), 1<<20) + for sc.Scan() { + var env workerctl.Envelope + Expect(json.Unmarshal(sc.Bytes(), &env)).To(Succeed(), sc.Text()) + out = append(out, env) + } + return out +} + +var _ = Describe("The runs of an agent worker", func() { + var ( + work *agentworker.Work + srv *httptest.Server + ) + + BeforeEach(func() { + work = agentworker.NewWork() + srv = httptest.NewServer(agentworker.Handler(agentworker.Config{}, work)) + DeferCleanup(srv.Close) + }) + + consume := func(kind messaging.WorkKind, max int, h messaging.WorkHandler) messaging.Subscription { + GinkgoHelper() + sub, err := work.Consume(GinkgoT().Context(), kind, max, h) + Expect(err).ToNot(HaveOccurred()) + return sub + } + + It("hands the payload and a publisher to the handler of the kind, and streams what it publishes", func() { + var payload []byte + consume(messaging.WorkAgentRun, 0, func(_ context.Context, p []byte, events messaging.Publisher) error { + payload = p + Expect(events.Publish("agent.helper.events.alice", map[string]string{"event_type": "json_message"})).To(Succeed()) + Expect(events.Publish("agent.helper.events.alice", map[string]string{"event_type": "json_message_status"})).To(Succeed()) + return nil + }) + resp := post(srv, workerctl.VerbAgentExecute, `{"agent_name":"helper"}`) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + Expect(resp.Header.Get("Content-Type")).To(Equal(workerctl.ContentTypeStream)) + lines := envelopes(resp) + Expect(lines).To(HaveLen(3)) + Expect(lines[0].Subject).To(Equal("agent.helper.events.alice")) + Expect(string(lines[0].Progress)).To(MatchJSON(`{"event_type":"json_message"}`)) + Expect(string(lines[1].Progress)).To(MatchJSON(`{"event_type":"json_message_status"}`)) + Expect(lines[2].Reply).ToNot(BeNil(), "the reply is the last line") + Expect(string(payload)).To(Equal(`{"agent_name":"helper"}`)) + }) + + It("takes the reply line of an MCP CI job from the result it published, and still streams that result", func() { + consume(messaging.WorkMCPCI, 1, func(_ context.Context, _ []byte, events messaging.Publisher) error { + Expect(events.Publish(messaging.SubjectJobProgress("j1"), map[string]string{"job_id": "j1", "status": "running"})).To(Succeed()) + Expect(events.Publish(messaging.SubjectJobResult("j1"), map[string]string{"job_id": "j1", "status": "completed", "result": "42"})).To(Succeed()) + return nil + }) + lines := envelopes(post(srv, workerctl.VerbMCPCIRun, `{"job_id":"j1"}`)) + Expect(lines).To(HaveLen(3)) + Expect(lines[1].Subject).To(Equal(messaging.SubjectJobResult("j1"))) + Expect(string(lines[2].Reply)).To(MatchJSON(`{"job_id":"j1","status":"completed","result":"42"}`)) + }) + + It("ends a run that published no result with an empty reply, which is still a reply", func() { + consume(messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + lines := envelopes(post(srv, workerctl.VerbAgentExecute, `{}`)) + Expect(lines).To(HaveLen(1)) + Expect(string(lines[0].Reply)).To(MatchJSON(`{}`)) + }) + + It("answers 500 and no reply when the handler cannot serve the run before it published anything", func() { + consume(messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { return errors.New("cannot serve") }) + resp := post(srv, workerctl.VerbAgentExecute, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError)) + Expect(readBody(resp)).To(ContainSubstring("cannot serve")) + }) + + It("ends the body with no reply line when the handler fails after it published, so the caller learns nothing", func() { + consume(messaging.WorkAgentRun, 0, func(_ context.Context, _ []byte, events messaging.Publisher) error { + _ = events.Publish("agent.a.events.u", map[string]string{"k": "v"}) + return errors.New("lost the plot") + }) + resp := post(srv, workerctl.VerbAgentExecute, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusOK)) + lines := envelopes(resp) + Expect(lines).To(HaveLen(1)) + Expect(lines[0].Reply).To(BeNil()) + }) + + It("survives a handler that panics, and ends that run with no reply", func() { + consume(messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { panic("bug") }) + resp := post(srv, workerctl.VerbAgentExecute, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError)) + consume(messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + Expect(post(srv, workerctl.VerbMCPCIRun, `{}`).StatusCode).To(Equal(http.StatusOK)) + }) + + It("refuses a body that cannot be a payload with 400 before the handler runs", func() { + // An empty body is no payload: the handler decodes the rest. + called := atomic.Bool{} + consume(messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { called.Store(true); return nil }) + resp, err := http.Get(srv.URL + workerctl.PathOf(workerctl.VerbAgentExecute)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusMethodNotAllowed)) + Expect(called.Load()).To(BeFalse()) + }) + + Describe("the limit on concurrent runs", func() { + It("answers busy when the limit is reached, and takes a run again once one ends", func() { + started := make(chan struct{}, 4) + release := make(chan struct{}) + consume(messaging.WorkMCPCI, 1, func(context.Context, []byte, messaging.Publisher) error { + started <- struct{}{} + <-release + return nil + }) + first := make(chan *http.Response, 1) + go func() { + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbMCPCIRun), "application/json", strings.NewReader(`{}`)) + Expect(err).ToNot(HaveOccurred()) + first <- resp + }() + Eventually(started).Should(Receive()) + + resp := post(srv, workerctl.VerbMCPCIRun, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + Expect(workerctl.IsBusy(resp.StatusCode, readBody(resp))).To(BeTrue()) + + close(release) + var done *http.Response + Eventually(first).Should(Receive(&done)) + _ = envelopes(done) + Expect(post(srv, workerctl.VerbMCPCIRun, `{}`).StatusCode).To(Equal(http.StatusOK)) + }) + + It("runs as many as asked at once, and any number for 0", func() { + for _, max := range []int{3, 0} { + work := agentworker.NewWork() + local := httptest.NewServer(agentworker.Handler(agentworker.Config{}, work)) + var running, peak atomic.Int32 + release := make(chan struct{}) + _, err := work.Consume(GinkgoT().Context(), messaging.WorkAgentRun, max, func(context.Context, []byte, messaging.Publisher) error { + n := running.Add(1) + for { + p := peak.Load() + if n <= p || peak.CompareAndSwap(p, n) { + break + } + } + <-release + running.Add(-1) + return nil + }) + Expect(err).ToNot(HaveOccurred()) + var wg sync.WaitGroup + for range 3 { + wg.Add(1) + go func() { + defer wg.Done() + resp, err := http.Post(local.URL+workerctl.PathOf(workerctl.VerbAgentExecute), "application/json", strings.NewReader(`{}`)) + if err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + }() + } + Eventually(running.Load, "5s").Should(BeEquivalentTo(3), "max=%d", max) + close(release) + wg.Wait() + local.Close() + } + }) + + It("answers busy for a kind that nothing consumes", func() { + resp := post(srv, workerctl.VerbAgentExecute, `{}`) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + Expect(workerctl.IsBusy(resp.StatusCode, readBody(resp))).To(BeTrue()) + }) + }) + + Describe("the context of a run", func() { + It("ends when the caller leaves", func() { + started := make(chan struct{}) + ended := make(chan error, 1) + consume(messaging.WorkAgentRun, 0, func(ctx context.Context, _ []byte, _ messaging.Publisher) error { + close(started) + <-ctx.Done() + ended <- ctx.Err() + return ctx.Err() + }) + ctx, cancel := context.WithCancel(context.Background()) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL+workerctl.PathOf(workerctl.VerbAgentExecute), bytes.NewReader([]byte(`{}`))) + go func() { _, _ = http.DefaultClient.Do(req) }() + Eventually(started).Should(BeClosed()) + cancel() + Eventually(ended, "5s").Should(Receive(MatchError(context.Canceled))) + }) + + It("ends when the worker stops consuming", func() { + started := make(chan struct{}) + ended := make(chan error, 1) + ctx, stop := context.WithCancel(context.Background()) + _, err := work.Consume(ctx, messaging.WorkAgentRun, 0, func(ctx context.Context, _ []byte, _ messaging.Publisher) error { + close(started) + <-ctx.Done() + ended <- ctx.Err() + return nil + }) + Expect(err).ToNot(HaveOccurred()) + go func() { + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbAgentExecute), "application/json", strings.NewReader(`{}`)) + if err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + }() + Eventually(started).Should(BeClosed()) + stop() + Eventually(ended, "5s").Should(Receive(MatchError(context.Canceled))) + }) + }) + + Describe("a subscription", func() { + It("refuses a kind that no verb serves", func() { + _, err := work.Consume(GinkgoT().Context(), messaging.WorkTask, 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + Expect(err).To(HaveOccurred()) + _, err = work.Consume(GinkgoT().Context(), "nonsense", 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + Expect(err).To(HaveOccurred()) + }) + + It("refuses a second consumer of the same kind", func() { + consume(messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + _, err := work.Consume(GinkgoT().Context(), messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + Expect(err).To(HaveOccurred()) + }) + + It("stops taking runs on Unsubscribe and waits for the runs in flight", func() { + started := make(chan struct{}) + release := make(chan struct{}) + sub := consume(messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { + close(started) + <-release + return nil + }) + go func() { + resp, err := http.Post(srv.URL+workerctl.PathOf(workerctl.VerbAgentExecute), "application/json", strings.NewReader(`{}`)) + if err == nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + }() + Eventually(started).Should(BeClosed()) + + unsubscribed := make(chan error, 1) + go func() { unsubscribed <- sub.Unsubscribe() }() + Consistently(unsubscribed, "100ms").ShouldNot(Receive(), "the run in flight is waited for") + Expect(post(srv, workerctl.VerbAgentExecute, `{}`).StatusCode).To(Equal(http.StatusServiceUnavailable), "no new run is taken") + close(release) + Eventually(unsubscribed).Should(Receive(Succeed())) + }) + }) +}) diff --git a/core/services/agentworker/start.go b/core/services/agentworker/start.go new file mode 100644 index 000000000000..e8f6df8575da --- /dev/null +++ b/core/services/agentworker/start.go @@ -0,0 +1,132 @@ +package agentworker + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "sync" + "time" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/worker" +) + +// Options is what Start needs to put an agent worker on the tunnel. +type Options struct { + // FrontendURL is the URL the worker registered with. + FrontendURL string + // NodeID is the identity that registration gave to the worker. + NodeID string + // TunnelToken returns the own tunnel credential of the node. A function, + // because the credential changes at every registration and the tunnel reads + // it at every dial. + TunnelToken func() string + // Reauthorize registers the node again after the frontend refused the + // credential. It may be nil. + Reauthorize func(ctx context.Context) error + // ControlToken guards the control server. It is the registration token of + // the deployment, which the control client of the frontend presents. An empty + // token leaves the server open to anything that can reach it, which is the + // tunnel and this host. + ControlToken string + // Handler serves the verbs. See Handler. + Handler http.Handler +} + +// Runtime is a running control plane: one HTTP server on loopback and one tunnel +// to the frontend. +type Runtime struct { + server *http.Server + tunnel *worker.Tunnel + addr string + + closeOnce sync.Once + closeErr error +} + +// Addr is the loopback address of the control server. It is never registered +// anywhere: the only way in is the tunnel. +func (r *Runtime) Addr() string { + if r == nil { + return "" + } + return r.addr +} + +// Connected reports whether the tunnel holds a session. It is about reach and +// nothing else, and nothing may read it as the worker being gone. +func (r *Runtime) Connected() bool { return r != nil && r.tunnel.Connected() } + +// Close stops the tunnel and the server. It is safe on nil and may be called +// more than once. +func (r *Runtime) Close() error { + if r == nil { + return nil + } + r.closeOnce.Do(func() { + r.closeErr = r.tunnel.Close() + shutdown(r.server) + }) + return r.closeErr +} + +// loopbackBind is where the server binds. A constant, so that "an agent worker +// opens no inbound port" is a fact about the code and not about its +// configuration: no flag moves it. +const loopbackBind = "127.0.0.1:0" + +// Start binds the control server on loopback, puts the bearer check in front of +// the handler, and holds a tunnel to the frontend that carries the requests. +// +// A tunnel that is refused to start is an error: the frontend URL or the +// identity of the node is unusable. A frontend that is down, or an admin who +// has not approved the node, is not: the tunnel retries with a backoff. +func Start(ctx context.Context, o Options) (*Runtime, error) { + if o.Handler == nil { + return nil, errors.New("starting the agent worker control plane: no handler") + } + if o.TunnelToken == nil { + return nil, errors.New("starting the agent worker control plane: no way to read the tunnel credential") + } + lis, err := net.Listen("tcp", loopbackBind) + if err != nil { + return nil, fmt.Errorf("binding the agent worker control server: %w", err) + } + addr := lis.Addr().String() + server := &http.Server{ + Handler: nodes.RequireBearer(o.ControlToken, o.Handler), + ReadHeaderTimeout: 10 * time.Second, + } + go func() { + if err := server.Serve(lis); err != nil && !errors.Is(err, http.ErrServerClosed) { + xlog.Error("The agent worker control server stopped", "error", err) + } + }() + + tun, err := worker.StartTunnel(ctx, worker.TunnelConfig{ + FrontendURL: o.FrontendURL, + NodeID: o.NodeID, + Token: o.TunnelToken, + Reauthorize: o.Reauthorize, + // An agent worker runs no backend process, so only the http tag is offered. + Services: worker.HTTPOnlyServices(addr), + }) + if err != nil { + shutdown(server) + return nil, fmt.Errorf("starting the agent worker tunnel: %w", err) + } + xlog.Info("Agent worker control plane serving over its tunnel", "node", o.NodeID) + return &Runtime{server: server, tunnel: tun, addr: addr}, nil +} + +func shutdown(server *http.Server) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := server.Shutdown(ctx); err != nil { + xlog.Debug("The agent worker control server did not stop cleanly", "error", err) + } +} diff --git a/core/services/agentworker/start_test.go b/core/services/agentworker/start_test.go new file mode 100644 index 000000000000..a0abf9ee0a9c --- /dev/null +++ b/core/services/agentworker/start_test.go @@ -0,0 +1,95 @@ +package agentworker_test + +import ( + "context" + "net" + "net/http" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/agentworker" + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +var _ = Describe("Starting the control plane of an agent worker", func() { + start := func(token string) *agentworker.Runtime { + GinkgoHelper() + rt, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{ + // Nothing listens here. A frontend that is away is not an error of + // Start: the tunnel retries with a backoff. + FrontendURL: "http://127.0.0.1:1", + NodeID: "node-1", + TunnelToken: func() string { return "tunnel-credential" }, + ControlToken: token, + Handler: agentworker.Handler(agentworker.Config{ + MCPTool: agentworker.Unary(func(context.Context, mcpremote.MCPToolRequest) mcpremote.MCPToolResponse { + return mcpremote.MCPToolResponse{Result: "ok"} + }), + }, nil), + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(rt.Close) + return rt + } + + call := func(rt *agentworker.Runtime, bearer string) *http.Response { + GinkgoHelper() + req, err := http.NewRequest(http.MethodPost, "http://"+rt.Addr()+workerctl.PathOf(workerctl.VerbMCPToolExecute), strings.NewReader(`{}`)) + Expect(err).ToNot(HaveOccurred()) + if bearer != "" { + req.Header.Set("Authorization", "Bearer "+bearer) + } + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + return resp + } + + It("binds loopback only, on a port it picks", func() { + rt := start("secret") + host, port, err := net.SplitHostPort(rt.Addr()) + Expect(err).ToNot(HaveOccurred()) + Expect(net.ParseIP(host).IsLoopback()).To(BeTrue()) + Expect(port).ToNot(Equal("0")) + }) + + It("serves a verb behind the bearer token and refuses a request without it", func() { + rt := start("secret") + Expect(call(rt, "secret").StatusCode).To(Equal(http.StatusOK)) + Expect(call(rt, "wrong").StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(call(rt, "").StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("reports a tunnel that is not connected as not connected, and does not fail to start", func() { + rt := start("secret") + Expect(rt.Connected()).To(BeFalse()) + }) + + It("stops serving on Close, and Close is safe on nil and twice", func() { + rt := start("secret") + addr := rt.Addr() + Expect(rt.Close()).To(Succeed()) + _, err := net.Dial("tcp", addr) + Expect(err).To(HaveOccurred()) + var none *agentworker.Runtime + Expect(none.Close()).To(Succeed()) + Expect(none.Connected()).To(BeFalse()) + }) + + It("refuses to start with no handler or no token function", func() { + _, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{FrontendURL: "http://127.0.0.1:1", NodeID: "n", TunnelToken: func() string { return "t" }}) + Expect(err).To(HaveOccurred()) + _, err = agentworker.Start(GinkgoT().Context(), agentworker.Options{FrontendURL: "http://127.0.0.1:1", NodeID: "n", Handler: http.NotFoundHandler()}) + Expect(err).To(HaveOccurred()) + }) + + It("refuses a frontend URL that cannot be dialled", func() { + _, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{ + FrontendURL: "ftp://frontend", NodeID: "n", TunnelToken: func() string { return "t" }, Handler: http.NotFoundHandler(), + }) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/agentworker/work.go b/core/services/agentworker/work.go new file mode 100644 index 000000000000..fae029a6e8ab --- /dev/null +++ b/core/services/agentworker/work.go @@ -0,0 +1,263 @@ +package agentworker + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "sync" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// Work is the WorkConsumer of an agent worker on the tunnel. A handler that was +// written for a queue of a bus runs here unchanged: Consume registers it, and +// the request that a frontend sends for a run is the delivery. +// +// The differences from a queue group are all in who chooses. A frontend picks the +// worker, so a worker that is full answers busy (see workerctl.WriteBusy) and the +// frontend offers the run to another worker. On a queue group a full worker held +// the delivery until it had room. +type Work struct { + mu sync.Mutex + runner map[string]*runner +} + +// NewWork returns a consumer with nothing registered. +func NewWork() *Work { return &Work{runner: map[string]*runner{}} } + +// verbOf maps a kind onto the verb that serves it. A plain task has none: no +// agent worker serves it. +func verbOf(kind messaging.WorkKind) (string, bool) { + switch kind { + case messaging.WorkAgentRun: + return workerctl.VerbAgentExecute, true + case messaging.WorkMCPCI: + return workerctl.VerbMCPCIRun, true + } + return "", false +} + +// runner is the registration of one kind. +type runner struct { + verb string + parent context.Context + handler messaging.WorkHandler + slots chan struct{} // nil: unbounded + + mu sync.Mutex + closed bool + running sync.WaitGroup +} + +// begin takes a slot, or reports that there is none. It never waits. +func (r *runner) begin() bool { + r.mu.Lock() + defer r.mu.Unlock() + if r.closed { + return false + } + if r.slots != nil { + select { + case r.slots <- struct{}{}: + default: + return false + } + } + r.running.Add(1) + return true +} + +func (r *runner) end() { + if r.slots != nil { + <-r.slots + } + r.running.Done() +} + +// Consume registers h for kind. maxInFlight bounds the runs at once: 0 and +// negative values are unbounded, as on the other carriers. Unsubscribe stops +// taking runs and waits for those in flight, so a handler must not call it. +func (w *Work) Consume(ctx context.Context, kind messaging.WorkKind, maxInFlight int, h messaging.WorkHandler) (messaging.Subscription, error) { + verb, ok := verbOf(kind) + if !ok { + return nil, fmt.Errorf("no agent worker verb serves the work kind %q", kind) + } + r := &runner{verb: verb, parent: ctx, handler: h} + if maxInFlight > 0 { + r.slots = make(chan struct{}, maxInFlight) + } + w.mu.Lock() + defer w.mu.Unlock() + if _, taken := w.runner[verb]; taken { + return nil, fmt.Errorf("the %s verb already has a consumer", verb) + } + w.runner[verb] = r + return &subscription{work: w, r: r}, nil +} + +type subscription struct { + work *Work + r *runner +} + +func (s *subscription) Unsubscribe() error { + s.r.mu.Lock() + s.r.closed = true + s.r.mu.Unlock() + s.work.mu.Lock() + if s.work.runner[s.r.verb] == s.r { + delete(s.work.runner, s.r.verb) + } + s.work.mu.Unlock() + s.r.running.Wait() + return nil +} + +func (w *Work) lookup(verb string) *runner { + w.mu.Lock() + defer w.mu.Unlock() + return w.runner[verb] +} + +// register mounts the verbs of the runs. A verb with no consumer answers busy and +// not 404: the worker serves it and has nobody to run it, which is a worker that +// is starting or stopping. +func (w *Work) register(mux *http.ServeMux) { + for _, kind := range []messaging.WorkKind{messaging.WorkAgentRun, messaging.WorkMCPCI} { + verb, _ := verbOf(kind) + mux.HandleFunc(workerctl.PathOf(verb), w.serve(verb)) + } +} + +func (w *Work) serve(verb string) http.HandlerFunc { + return func(rw http.ResponseWriter, req *http.Request) { + body, ok := workerctl.ReadRequestBody(rw, req) + if !ok { + return + } + r := w.lookup(verb) + if r == nil || !r.begin() { + workerctl.WriteBusy(rw) + return + } + defer r.end() + + // The run ends when the caller leaves and when the worker stops: either + // way nobody is left to read its result. + ctx, cancel := context.WithCancel(req.Context()) + defer cancel() + stop := context.AfterFunc(r.parent, cancel) + defer stop() + + stream := &stream{w: rw} + err := runHandler(ctx, r.handler, body, stream) + if err != nil { + if !stream.committed() { + failToServe(rw, verb, err) + return + } + // The status is spent. A body with no reply line is a link that broke, + // which is the only honest thing to tell the caller: this worker did + // not answer. + xlog.Warn("An agent worker run failed after it had streamed events, so its response ends with no reply line", + "verb", verb, "error", err) + return + } + stream.finish(verb) + } +} + +// runHandler calls h and turns a panic into an error, so one bad run does not +// take the worker down with it. +func runHandler(ctx context.Context, h messaging.WorkHandler, payload []byte, events messaging.Publisher) (err error) { + defer func() { + if p := recover(); p != nil { + err = fmt.Errorf("the run panicked: %v", p) + } + }() + return h(ctx, payload, events) +} + +// stream writes the lines of one streaming response. It is the Publisher that a +// handler gets as the place for its events. +// +// A handler may publish from several goroutines, so the mutex is not optional: +// without it a progress line could interleave with the reply and put a torn line +// on the wire. done keeps the reply the last line. +type stream struct { + mu sync.Mutex + w http.ResponseWriter + header bool + done bool + result *workerctl.RunReply +} + +// Publish writes one event as a progress line that names its subject. The result +// of a job is also kept: it becomes the reply line. +func (s *stream) Publish(subject string, data any) error { + raw, err := json.Marshal(data) + if err != nil { + return fmt.Errorf("encoding an event: %w", err) + } + s.mu.Lock() + defer s.mu.Unlock() + if s.done { + return nil + } + if messaging.SubjectMatches(messaging.SubjectJobResultWildcard, subject) { + var res workerctl.RunReply + if json.Unmarshal(raw, &res) == nil { + s.result = &res + } + } + s.write(workerctl.Envelope{Subject: subject, Progress: raw}) + return nil +} + +func (s *stream) committed() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.header +} + +// finish writes the reply line: the result the run published, or an empty reply +// for a run that publishes none. +func (s *stream) finish(verb string) { + s.mu.Lock() + defer s.mu.Unlock() + if s.done { + return + } + reply := workerctl.RunReply{} + if s.result != nil { + reply = *s.result + } + raw, err := json.Marshal(reply) + if err != nil { + raw = json.RawMessage(`{}`) + } + s.write(workerctl.Envelope{Reply: raw}) + s.done = true + xlog.Debug("agent worker run finished", "verb", verb) +} + +// write encodes one line and flushes it. The caller holds s.mu. +func (s *stream) write(env workerctl.Envelope) { + if !s.header { + s.w.Header().Set("Content-Type", workerctl.ContentTypeStream) + s.w.Header().Set("X-Content-Type-Options", "nosniff") + s.w.WriteHeader(http.StatusOK) + s.header = true + } + if err := json.NewEncoder(s.w).Encode(env); err != nil { + xlog.Debug("agent worker stream line could not be written", "error", err) + return + } + if f, ok := s.w.(http.Flusher); ok { + f.Flush() + } +} diff --git a/core/services/nodes/file_transfer_server.go b/core/services/nodes/file_transfer_server.go index 11e4c9c23b88..97fce01921f1 100644 --- a/core/services/nodes/file_transfer_server.go +++ b/core/services/nodes/file_transfer_server.go @@ -1095,6 +1095,19 @@ func checkBearerToken(r *http.Request, expectedToken string) bool { return subtle.ConstantTimeCompare([]byte(provided), []byte(expectedToken)) == 1 } +// RequireBearer puts the bearer check of the worker in front of next. A request +// without the token is answered 401. An empty token lets everything through, as +// it does for the file routes. +func RequireBearer(token string, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !checkBearerToken(r, token) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + next.ServeHTTP(w, r) + }) +} + // validatePathInDir checks that targetPath is within the given base directory. func validatePathInDir(targetPath, baseDir string) error { absBase, err := filepath.Abs(baseDir) diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index 7542d16de4b9..9e3ade05f29a 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -930,6 +930,16 @@ func tunnelServices(cfg *Config, httpBindAddr string) map[string]LocalService { } } +// HTTPOnlyServices is the routing table of a worker that runs no backend +// processes, such as an agent worker. It offers the http tag, which reaches the +// server at httpBindAddr, and no gRPC tag, so the frontend has no way to name a +// port of this host. +func HTTPOnlyServices(httpBindAddr string) map[string]LocalService { + return map[string]LocalService{ + tunnel.StreamTagHTTP: fixedService(loopbackAddr(httpBindAddr)), + } +} + // fixedService routes a tagged stream to one address of this worker and ignores // what the frontend named. A worker has one HTTP server and only the worker // knows where it listens, so the frontend has nothing useful to say about the diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go index b04d99e24dd7..53d698e0c739 100644 --- a/core/services/worker/tunnel_test.go +++ b/core/services/worker/tunnel_test.go @@ -1081,6 +1081,24 @@ var _ = Describe("Worker tunnel local services", func() { }) }) + Describe("HTTPOnlyServices, the table of a worker that runs no backend", func() { + It("serves the http tag and refuses to dial a backend", func() { + table := HTTPOnlyServices("0.0.0.0:50050") + Expect(table).To(HaveLen(1)) + Expect(table).To(HaveKey(tunnel.StreamTagHTTP)) + Expect(table).ToNot(HaveKey(tunnel.StreamTagGRPC)) + }) + + It("dials its own server whatever address the frontend names", func() { + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = ln.Close() }() + conn, err := HTTPOnlyServices(ln.Addr().String())[tunnel.StreamTagHTTP](ctx, "attacker.invalid:1") + Expect(err).ToNot(HaveOccurred()) + _ = conn.Close() + }) + }) + DescribeTable("tunnelEndpoint builds the URL that the worker dials", func(frontend, nodeID, want string) { got, err := tunnelEndpoint(frontend, nodeID) From 294b08432fb46652d5b7b458991f5be101a2f111 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:05:30 +0000 Subject: [PATCH 077/162] feat(cli): let the agent worker attach to the tunnel carrier The agent worker refused a cluster on the tunnel because it could serve agent runs over NATS only. It now attaches to the carrier that the registration names. On the tunnel it opens no connection to a bus: it registers its consumers, serves the MCP requests, the agent runs and the MCP CI jobs on a loopback server that the frontend reaches through the tunnel of the worker, and reads its tunnel credential at every dial. The NATS path is the code that was there. A frontend that names a carrier the worker does not know gets a message that asks for an upgrade. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/agent_worker.go | 130 ++++++++++++++++++++++---- core/cli/agent_worker_subject_test.go | 31 ++++-- 2 files changed, 132 insertions(+), 29 deletions(-) diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 9ee8ecc57928..3401cd6cd278 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -16,10 +16,12 @@ import ( "github.com/mudler/LocalAI/core/config" mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/agentworker" "github.com/mudler/LocalAI/core/services/jobs" mcpRemote "github.com/mudler/LocalAI/core/services/mcp" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/cogito" @@ -28,17 +30,19 @@ import ( ) // AgentWorkerCMD starts a dedicated agent worker process for distributed mode. -// It registers with the frontend, subscribes to the NATS agent execution queue, -// and executes agent chats using cogito. The worker is a pure executor — it -// receives the full agent config and skills in the NATS job payload, so it -// does not need direct database access. +// It registers with the frontend, attaches to the carrier that the frontend +// names, and executes agent chats using cogito. On NATS it subscribes to the +// agent execution queue. On the tunnel it holds one outbound connection to the +// frontend and serves the requests that arrive on it. The worker is a pure +// executor — it receives the full agent config and skills in the job payload, +// so it does not need direct database access. // // Usage: // // localai agent-worker --nats-url nats://... --register-to http://localai:8080 type AgentWorkerCMD struct { // NATS (required) - NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed while the cluster runs on NATS" group:"distributed"` + NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed only while the cluster runs on NATS" group:"distributed"` // Registration (required) RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"` @@ -88,20 +92,25 @@ func validateAgentSubject(subject string) error { return nil } -// checkAgentCarrier checks that the carrier which the frontend named is one that -// this agent worker can use. It serves its agent runs over NATS only. The flag of -// the NATS URL cannot be required, because the carrier is known only after the -// registration, so the check is here. A frontend that names no carrier predates -// carriers and runs on NATS. -func checkAgentCarrier(carrier, natsURL string) error { - switch carrier { +// agentCarrier maps the carrier that the frontend named onto the one this agent +// worker attaches to, and reports whether it is the tunnel. A frontend that names +// none predates carriers and runs on NATS. The flag of the NATS URL cannot be +// required, because the carrier is known only after the registration, so the +// check is here. +func agentCarrier(named, natsURL string) (onTunnel bool, err error) { + switch named { + case "tunnel": + if natsURL != "" { + xlog.Info("The cluster runs on the tunnel, so the NATS URL of this agent worker is not used") + } + return true, nil case "", "nats": if natsURL == "" { - return fmt.Errorf("the cluster runs on NATS and this agent worker has no NATS URL: set LOCALAI_NATS_URL") + return false, fmt.Errorf("the cluster runs on NATS and this agent worker has no NATS URL: set LOCALAI_NATS_URL") } - return nil + return false, nil default: - return fmt.Errorf("the cluster carrier is %q, which this agent worker cannot use: it serves agent runs over NATS only", carrier) + return false, fmt.Errorf("the frontend names the carrier %q, which this agent worker does not know: upgrade the agent worker", named) } } @@ -156,7 +165,8 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { nodeID := res.ID xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo) - if err := checkAgentCarrier(res.Carrier, cmd.NatsURL); err != nil { + onTunnel, err := agentCarrier(res.Carrier, cmd.NatsURL) + if err != nil { return err } @@ -174,6 +184,10 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { go regClient.HeartbeatLoop(shutdownCtx, nodeID, heartbeatInterval, func() map[string]any { return map[string]any{} }) + if onTunnel { + return cmd.runOnTunnel(shutdownCtx, shutdownCancel, regClient, credMgr, nodeID, apiURL) + } + // Resolve NATS credentials with precedence: explicit env override, then // frontend-minted (auto-refreshed before expiry), then service fallback. // Each static source must supply JWT and seed together. @@ -281,19 +295,95 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // Wait for an OS signal or an internal fatal condition (e.g. NATS // credentials became unrenewable), so the worker restarts and re-acquires // rather than lingering unable to serve. + runErr := awaitShutdown(shutdownCtx, "NATS credentials unavailable") + + xlog.Info("Shutting down agent worker") + shutdownCancel() // stop heartbeat loop immediately + dispatcher.Stop() + mcpTools.CloseAllMCPSessions() + regClient.GracefulDeregister(nodeID) + return runErr +} + +// awaitShutdown blocks until the process is told to stop or the context of the +// worker ends for a fatal reason, and returns an error in the second case so +// that the worker restarts and registers again instead of lingering unable to +// serve. +func awaitShutdown(shutdownCtx context.Context, fatal string) error { sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) - var runErr error + defer signal.Stop(sigCh) select { case <-sigCh: + return nil case <-shutdownCtx.Done(): - runErr = fmt.Errorf("agent worker shutting down: NATS credentials unavailable") - xlog.Error("Internal shutdown requested", "error", runErr) + err := fmt.Errorf("agent worker shutting down: %s", fatal) + xlog.Error("Internal shutdown requested", "error", err) + return err } +} +// runOnTunnel serves the agent worker on the tunnel carrier: no connection to a +// bus, one HTTP server on loopback that the frontend reaches through the tunnel +// of this worker, and the same handlers as on NATS. The frontend picks the worker +// and sends the request, so the dispatcher below gets its deliveries from a +// consumer that the control server feeds. +func (cmd *AgentWorkerCMD) runOnTunnel(shutdownCtx context.Context, shutdownCancel context.CancelFunc, regClient *workerregistry.RegistrationClient, credMgr *workerregistry.CredentialManager, nodeID, apiURL string) error { + work := agentworker.NewWork() + + // No bus, so no broadcaster: the events of a run are published on the + // publisher of its delivery, which writes them to the frontend that drives the + // run. The cancel of a run is its request: the frontend ends the stream. + eventBridge := agents.NewEventBridge(nil, nil, "agent-worker-"+nodeID) + dispatcher := agents.NewNATSDispatcher(work, eventBridge, nil, apiURL, cmd.APIToken, 0) + // The consumers are registered before the tunnel starts, so that a frontend + // which connects at once finds a worker that can take its first run. + if err := dispatcher.Start(shutdownCtx); err != nil { + return fmt.Errorf("starting dispatcher: %w", err) + } + + mcpCIJobTimeout, err := time.ParseDuration(cmd.MCPCIJobTimeout) + if err != nil && cmd.MCPCIJobTimeout != "" { + xlog.Warn("invalid MCP CI job timeout, using default 10m", "input", cmd.MCPCIJobTimeout, "error", err) + } + mcpCIJobTimeout = cmp.Or(mcpCIJobTimeout, config.DefaultMCPCIJobTimeout) + if _, err := startMCPCIConsumer(shutdownCtx, work, apiURL, cmd.APIToken, mcpCIJobTimeout); err != nil { + return err + } + + runtime, err := agentworker.Start(shutdownCtx, agentworker.Options{ + FrontendURL: cmd.RegisterTo, + NodeID: nodeID, + // The credential is new after every registration, so it is read at every + // dial and not once. + TunnelToken: credMgr.TunnelToken, + Reauthorize: credMgr.Reregister, + ControlToken: cmd.RegistrationToken, + Handler: agentworker.Handler(agentworker.Config{ + MCPTool: agentworker.Unary(handleMCPToolRequest), + MCPDiscovery: agentworker.Unary(handleMCPDiscoveryRequest), + // The node's backend stop: the agent worker has no model loader, so it + // drops the sessions cached for the backend that goes away. + BackendStop: func(_ context.Context, req workerctl.BackendStopRequest) error { + if req.Backend != "" { + mcpTools.CloseMCPSessions(req.Backend) + } + return nil + }, + }, work), + }) + if err != nil { + return err + } + xlog.Info("Agent worker ready, waiting for runs over its tunnel") + + runErr := awaitShutdown(shutdownCtx, "the control plane stopped") xlog.Info("Shutting down agent worker") - shutdownCancel() // stop heartbeat loop immediately + shutdownCancel() dispatcher.Stop() + if err := runtime.Close(); err != nil { + xlog.Warn("Closing the agent worker tunnel failed", "error", err) + } mcpTools.CloseAllMCPSessions() regClient.GracefulDeregister(nodeID) return runErr diff --git a/core/cli/agent_worker_subject_test.go b/core/cli/agent_worker_subject_test.go index 9b4537407743..6bf175c461aa 100644 --- a/core/cli/agent_worker_subject_test.go +++ b/core/cli/agent_worker_subject_test.go @@ -30,19 +30,32 @@ var _ = Describe("validateAgentSubject", func() { }) }) -var _ = Describe("checkAgentCarrier", func() { - It("accepts NATS with a URL, named or not named", func() { - Expect(checkAgentCarrier("nats", "nats://n:4222")).To(Succeed()) - Expect(checkAgentCarrier("", "nats://n:4222")).To(Succeed()) +var _ = Describe("agentCarrier", func() { + It("attaches to NATS when the frontend names it or names nothing, and the URL is there", func() { + for _, named := range []string{"nats", ""} { + onTunnel, err := agentCarrier(named, "nats://n:4222") + Expect(err).ToNot(HaveOccurred(), named) + Expect(onTunnel).To(BeFalse(), named) + } }) It("asks for the URL when the cluster runs on NATS and there is none", func() { - Expect(checkAgentCarrier("nats", "")).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) - Expect(checkAgentCarrier("", "")).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) + _, err := agentCarrier("nats", "") + Expect(err).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) + _, err = agentCarrier("", "") + Expect(err).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) }) - It("says that it cannot use the tunnel, with or without a URL", func() { - Expect(checkAgentCarrier("tunnel", "")).To(MatchError(ContainSubstring("cannot use"))) - Expect(checkAgentCarrier("tunnel", "nats://n:4222")).To(MatchError(ContainSubstring("cannot use"))) + It("attaches to the tunnel when the frontend names it, with a URL or without one", func() { + for _, url := range []string{"", "nats://n:4222"} { + onTunnel, err := agentCarrier("tunnel", url) + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(BeTrue()) + } + }) + + It("asks for an upgrade when the frontend names a carrier it does not know", func() { + _, err := agentCarrier("smoke-signals", "nats://n:4222") + Expect(err).To(MatchError(ContainSubstring("upgrade"))) }) }) From 48f9a99cfcf85c2c2460f283845eae77dfbbf14a Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:08:53 +0000 Subject: [PATCH 078/162] test(messaging): add one conformance suite for the carriers of the work queue The queue of the NATS carrier had specs of its own and nothing that another carrier could be held to. RunWorkQueueConformance states what every carrier of WorkQueue and WorkConsumer shares: one consumer gets the payload and a publisher for its events, competing consumers split the work once when nothing fails, the kinds stay apart, the limit on concurrent handlers holds, Unsubscribe waits for the handler that runs, and one payload bound applies. It leaves out what differs and what the contract leaves to the carrier: the fate of work whose consumer dies, and how often a handler that failed is called again. A handler must tolerate a repeat. MaxWorkPayloadBytes is that bound, 1 MiB, the default max_payload of NATS, so that a job one carrier accepts is not refused by another. The NATS queue passes the suite. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../messaging/client_conformance_test.go | 40 +++ core/services/messaging/limits.go | 19 ++ .../messaging/messagingtest/workqueue.go | 244 ++++++++++++++++++ 3 files changed, 303 insertions(+) create mode 100644 core/services/messaging/messagingtest/workqueue.go diff --git a/core/services/messaging/client_conformance_test.go b/core/services/messaging/client_conformance_test.go index e08d7171de84..572afdf884bf 100644 --- a/core/services/messaging/client_conformance_test.go +++ b/core/services/messaging/client_conformance_test.go @@ -88,3 +88,43 @@ var _ = Describe("NATS client", func() { } }) }) + +// The queue groups of NATS run the shared work queue suite. Each consumer has a +// client of its own, as each worker does, so that they compete for the queue. +var _ = Describe("NATS work queue", func() { + messagingtest.RunWorkQueueConformance(func() messagingtest.WorkQueueRig { + url, err := sharedNATS() + if err != nil { + if os.Getenv("CI") != "" && runtime.GOOS != "darwin" { + Fail("testcontainers requires Docker and CI is set: " + err.Error()) + } + Skip("testcontainers requires Docker: " + err.Error()) + } + producer, err := messaging.New(url) + if err != nil { + Fail("connecting to the test NATS server: " + err.Error()) + } + var mu sync.Mutex + clients := []*messaging.Client{producer} + return messagingtest.WorkQueueRig{ + Queue: messaging.NewNATSWorkQueue(producer), + NewConsumer: func() messaging.WorkConsumer { + c, err := messaging.New(url) + if err != nil { + Fail("connecting to the test NATS server: " + err.Error()) + } + mu.Lock() + clients = append(clients, c) + mu.Unlock() + return messaging.NewNATSWorkConsumer(c) + }, + Cleanup: func() { + mu.Lock() + defer mu.Unlock() + for _, c := range clients { + c.Close() + } + }, + } + }) +}) diff --git a/core/services/messaging/limits.go b/core/services/messaging/limits.go index 7f52f1416ef0..941142f90345 100644 --- a/core/services/messaging/limits.go +++ b/core/services/messaging/limits.go @@ -26,3 +26,22 @@ func CheckBroadcastSize(subject string, size int) error { } return nil } + +// MaxWorkPayloadBytes is the largest encoded payload that a WorkQueue accepts. +// Every carrier of the queue enforces it, so a job that one carrier takes is not +// refused by another after a change of carrier. NATS refuses a payload above the +// max_payload of its server, which is 1 MiB unless an operator raised it, and the +// bound is that default, for the carriers that have no limit of their own. +const MaxWorkPayloadBytes = 1 << 20 + +// ErrWorkPayloadTooLarge means a payload is larger than MaxWorkPayloadBytes. +var ErrWorkPayloadTooLarge = errors.New("messaging: work payload is too large") + +// CheckWorkSize returns an error that wraps ErrWorkPayloadTooLarge when an +// encoded payload of size bytes is above MaxWorkPayloadBytes. +func CheckWorkSize(kind WorkKind, size int) error { + if size > MaxWorkPayloadBytes { + return fmt.Errorf("%w: %d bytes for kind %q, the bound is %d", ErrWorkPayloadTooLarge, size, kind, MaxWorkPayloadBytes) + } + return nil +} diff --git a/core/services/messaging/messagingtest/workqueue.go b/core/services/messaging/messagingtest/workqueue.go new file mode 100644 index 000000000000..2b1d2b767c5a --- /dev/null +++ b/core/services/messaging/messagingtest/workqueue.go @@ -0,0 +1,244 @@ +package messagingtest + +import ( + "context" + "encoding/json" + "strings" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// WorkQueueRig is what a factory returns for the work queue suite: one queue, +// with a producer and as many competing consumers as the suite asks for. +type WorkQueueRig struct { + // Queue is the producer. + Queue messaging.WorkQueue + // NewConsumer returns a consumer that competes with every other consumer of + // the rig for the work of Queue, as one more worker would. + NewConsumer func() messaging.WorkConsumer + // Cleanup releases the rig. + Cleanup func() +} + +// WorkQueueFactory returns a ready rig. A factory that cannot build one calls +// Skip or Fail itself. +type WorkQueueFactory func() WorkQueueRig + +type job struct { + ID string `json:"id"` + Data string `json:"data,omitempty"` +} + +// RunWorkQueueConformance registers the suite every WorkQueue carrier must pass. +// +// It states what the carriers share: one consumer receives a payload, competing +// consumers split the work, the kinds are apart, the limit on concurrent +// handlers holds, and the same payload bound applies. It does not state what +// differs and what the contract leaves to the carrier: what happens to a unit of +// work whose consumer dies, and how often a handler that failed is called again. +// Handlers must tolerate a repeat. +func RunWorkQueueConformance(newRig WorkQueueFactory) { + Describe("WorkQueue conformance", func() { + var ( + rig WorkQueueRig + ctx context.Context + ) + + BeforeEach(func() { + rig = newRig() + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(context.Background()) + DeferCleanup(cancel) + }) + // A factory that skips (no Docker for the NATS run) returns before it sets + // a cleanup, so the guard keeps a skip from turning into a panic. + AfterEach(func() { + if rig.Cleanup != nil { + rig.Cleanup() + } + }) + + // seen keeps the ids a handler has been given. + type seen struct { + mu sync.Mutex + ids []string + } + add := func(s *seen, id string) { + s.mu.Lock() + defer s.mu.Unlock() + s.ids = append(s.ids, id) + } + count := func(s *seen) func() int { + return func() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.ids) + } + } + + consume := func(c messaging.WorkConsumer, kind messaging.WorkKind, max int, h messaging.WorkHandler) messaging.Subscription { + GinkgoHelper() + sub, err := c.Consume(ctx, kind, max, h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + return sub + } + + It("gives the payload to a consumer of the kind, with a publisher for its events", func() { + var got seen + var payload atomic.Value + var hadEvents atomic.Bool + consume(rig.NewConsumer(), messaging.WorkMCPCI, 0, func(_ context.Context, p []byte, events messaging.Publisher) error { + payload.Store(append([]byte(nil), p...)) + hadEvents.Store(events != nil) + add(&got, "x") + return nil + }) + + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, job{ID: "j1", Data: "hello"})).To(Succeed()) + + Eventually(count(&got), 10*time.Second).Should(Equal(1)) + Expect(payload.Load().([]byte)).To(MatchJSON(`{"id":"j1","data":"hello"}`)) + Expect(hadEvents.Load()).To(BeTrue()) + }) + + It("gives each unit of work to one of the competing consumers, once, when nothing fails", func() { + const n = 60 + var got seen + for range 3 { + consume(rig.NewConsumer(), messaging.WorkAgentRun, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + var j job + if err := json.Unmarshal(p, &j); err != nil { + return nil + } + add(&got, j.ID) + return nil + }) + } + for i := range n { + Expect(rig.Queue.Enqueue(ctx, messaging.WorkAgentRun, job{ID: "j" + string(rune('A'+i/26)) + string(rune('a'+i%26))})).To(Succeed()) + } + + Eventually(count(&got), 20*time.Second).Should(Equal(n)) + Consistently(count(&got), 500*time.Millisecond).Should(Equal(n), "no unit is given twice") + got.mu.Lock() + defer got.mu.Unlock() + unique := map[string]bool{} + for _, id := range got.ids { + unique[id] = true + } + Expect(unique).To(HaveLen(n)) + }) + + It("keeps the kinds apart", func() { + var agent, ci seen + consume(rig.NewConsumer(), messaging.WorkAgentRun, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + add(&agent, string(p)) + return nil + }) + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, job{ID: "ci-1"})).To(Succeed()) + Consistently(count(&agent), 500*time.Millisecond).Should(BeZero()) + + consume(rig.NewConsumer(), messaging.WorkMCPCI, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + add(&ci, string(p)) + return nil + }) + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, job{ID: "ci-2"})).To(Succeed()) + Eventually(count(&ci), 10*time.Second).Should(BeNumerically(">=", 1)) + Expect(count(&agent)()).To(BeZero()) + }) + + It("refuses a kind that it does not know", func() { + err := rig.Queue.Enqueue(ctx, messaging.WorkKind("nonsense"), job{ID: "x"}) + Expect(err).To(HaveOccurred()) + _, err = rig.NewConsumer().Consume(ctx, messaging.WorkKind("nonsense"), 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + Expect(err).To(HaveOccurred()) + }) + + DescribeTable("runs no more handlers at once than it was asked for", + func(limit, items int) { + var running, peak, finished atomic.Int32 + release := make(chan struct{}) + consume(rig.NewConsumer(), messaging.WorkMCPCI, limit, func(context.Context, []byte, messaging.Publisher) error { + n := running.Add(1) + for { + p := peak.Load() + if n <= p || peak.CompareAndSwap(p, n) { + break + } + } + <-release + running.Add(-1) + finished.Add(1) + return nil + }) + for i := range items { + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, job{ID: "j" + string(rune('a'+i))})).To(Succeed()) + } + + Eventually(running.Load, 10*time.Second).Should(BeEquivalentTo(limit)) + Consistently(running.Load, 500*time.Millisecond).Should(BeEquivalentTo(limit)) + close(release) + Eventually(finished.Load, 10*time.Second).Should(BeEquivalentTo(items)) + Expect(peak.Load()).To(BeEquivalentTo(limit)) + }, + Entry("one at a time", 1, 4), + Entry("three at a time", 3, 7), + ) + + It("lets Unsubscribe wait for the handler that runs, and takes no more work after it", func() { + started := make(chan struct{}) + release := make(chan struct{}) + var after seen + sub, err := rig.NewConsumer().Consume(ctx, messaging.WorkMCPCI, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + var j job + _ = json.Unmarshal(p, &j) + if j.ID == "first" { + close(started) + <-release + return nil + } + add(&after, j.ID) + return nil + }) + Expect(err).ToNot(HaveOccurred()) + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, job{ID: "first"})).To(Succeed()) + Eventually(started, 10*time.Second).Should(BeClosed()) + + done := make(chan struct{}) + go func() { defer close(done); _ = sub.Unsubscribe() }() + Consistently(done, 200*time.Millisecond).ShouldNot(BeClosed(), "Unsubscribe waits for the running handler") + close(release) + Eventually(done, 5*time.Second).Should(BeClosed()) + + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, job{ID: "second"})).To(Succeed()) + Consistently(count(&after), 500*time.Millisecond).Should(BeZero()) + }) + + It("accepts a payload up to the bound and refuses one above it", func() { + var got atomic.Int32 + var size atomic.Int32 + consume(rig.NewConsumer(), messaging.WorkMCPCI, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + size.Store(int32(len(p))) + got.Add(1) + return nil + }) + // The envelope of the job is a few bytes, so a margin keeps the + // accepted payload under the bound after encoding. + big := job{ID: "big", Data: strings.Repeat("a", messaging.MaxWorkPayloadBytes-1024)} + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, big)).To(Succeed()) + Eventually(got.Load, 15*time.Second).Should(BeEquivalentTo(1)) + Expect(int(size.Load())).To(BeNumerically(">", messaging.MaxWorkPayloadBytes-2048)) + + tooBig := job{ID: "too-big", Data: strings.Repeat("a", messaging.MaxWorkPayloadBytes+1)} + Expect(rig.Queue.Enqueue(ctx, messaging.WorkMCPCI, tooBig)).ToNot(Succeed()) + Consistently(got.Load, 500*time.Millisecond).Should(BeEquivalentTo(1)) + }) + }) +} From e97729b2cb7d6e9500b9712354cef4fe0f8764b2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:11:30 +0000 Subject: [PATCH 079/162] feat(jobs): add the claim table, the queue of the carrier that has no broker A queue group delivers a message to one of several competing consumers. A row and a lock do the same, and the row stays until the work has an answer. The claim table keeps a row in one of three states: pending, claimed by a frontend replica, or migrated out by a change of carrier. ClaimNext takes the oldest pending row of a kind with FOR UPDATE SKIP LOCKED in one statement, so competing replicas never block each other and never take one row twice. ReleaseClaim sends a row back for a failure that says nothing about the work, with a wait that doubles from two seconds to a minute, so that one poison row cannot be claimed ahead of every newer row on every tick. CompleteClaim deletes the row once the work has an answer. A claim is settled only by the replica and the attempt that hold it. A holder that was reaped and answers late finds the row changed and does nothing, so it cannot delete the claim of the replica that took the work over. ReapAbandoned returns the claims of replicas that are not live. It reads the same predicate of liveness as the rest of the cluster and has no age in it: a replica that heartbeats and holds a claim for an hour has a slow job, and taking it would run the work twice. MigratePending takes the pending rows out for a carrier that has its own queue, once. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/jobs/claim.go | 334 +++++++++++++++++++++ core/services/jobs/claim_test.go | 493 +++++++++++++++++++++++++++++++ 2 files changed, 827 insertions(+) create mode 100644 core/services/jobs/claim.go create mode 100644 core/services/jobs/claim_test.go diff --git a/core/services/jobs/claim.go b/core/services/jobs/claim.go new file mode 100644 index 000000000000..870dd0eaf857 --- /dev/null +++ b/core/services/jobs/claim.go @@ -0,0 +1,334 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "github.com/google/uuid" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/advisorylock" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" +) + +// The claim queue is the WorkQueue of the carrier that has no broker. A queue +// group did one thing: it delivered a message to one of several competing +// consumers. A row and a lock do the same, and the row stays until someone has +// done the work, which a published message did not. +// +// The states of a row: +// +// pending -> claimed a replica took it (ClaimNext) +// claimed -> pending the work got no answer (ReleaseClaim), or the replica +// that held it is gone (ReapAbandoned) +// claimed -> deleted the work was answered (CompleteClaim) +// pending -> migrated a change of carrier took it out of the queue (MigratePending) +const ( + ClaimPending = "pending" + ClaimClaimed = "claimed" + ClaimMigrated = "migrated" +) + +// WorkClaim is one unit of queued work. +// +// ClaimedBy names the frontend replica that holds the claim, never the worker +// that the work was handed to. A worker that dies in the middle of a run breaks +// the stream that the replica is reading, and the replica settles the claim. It +// is the death of the replica that leaves a row nobody drives, and only the reap +// answers that. +type WorkClaim struct { + ID string `gorm:"primaryKey;size:36"` + Kind string `gorm:"size:32;index:idx_work_claims_pick,priority:1"` + State string `gorm:"size:16;not null;default:pending;index:idx_work_claims_pick,priority:2"` + Payload []byte + ClaimedBy string `gorm:"size:64;index"` + ClaimedAt *time.Time + // Attempts counts the times the row went back to the pool. It also fences: + // a holder that was reaped and comes back with its answer carries a count + // that no longer matches the row, so it cannot settle the claim of the + // replica that took the work over. + Attempts int + // NotBefore is the earliest the row may be claimed again. The database clock + // stamps it. NULL means now. + NotBefore *time.Time `gorm:"index"` + CreatedAt time.Time `gorm:"index:idx_work_claims_pick,priority:3"` +} + +// TableName pins the table, because the statements below name it in raw SQL and a +// change of the naming rule of gorm would leave the two spellings disagreeing at +// run time and not at build time. +func (WorkClaim) TableName() string { return claimsTable } + +const claimsTable = "work_claims" + +// ErrNoWork reports that no row of the kind was there to claim. It is an +// ordinary emptiness and never an error about the database. +var ErrNoWork = errors.New("jobs: no unclaimed work") + +// MigrateClaims creates the claim table, under the lock that every schema +// migration of the deployment takes, because frontends start together and two +// concurrent AutoMigrate calls on one table race in the catalog. +func MigrateClaims(ctx context.Context, db *gorm.DB) error { + if err := advisorylock.WithLockCtx(ctx, db, advisorylock.KeySchemaMigrate, func() error { + return db.WithContext(ctx).AutoMigrate(&WorkClaim{}) + }); err != nil { + return fmt.Errorf("migrating the claim table: %w", err) + } + return nil +} + +// requirePostgres refuses an operation whose statement only PostgreSQL runs. The +// refusal matters more than the syntax error it replaces: SKIP LOCKED and +// make_interval fail on SQLite with a parse error that reads like a missing +// migration, and a queue that cannot claim must not be mistaken for one whose +// table is not there yet. +func requirePostgres(db *gorm.DB, op string) error { + if db == nil { + return fmt.Errorf("%s: no database handle", op) + } + if name := db.Dialector.Name(); !strings.Contains(name, "postgres") { + return fmt.Errorf("%s: the claim queue requires PostgreSQL, this deployment runs on %q", op, name) + } + return nil +} + +// knownKind reports whether kind is one of the work kinds. A row of another kind +// would be one that no consumer ever asks for. +func knownKind(kind messaging.WorkKind) bool { + switch kind { + case messaging.WorkTask, messaging.WorkMCPCI, messaging.WorkAgentRun: + return true + } + return false +} + +// EnqueueClaim writes one pending row and returns its id. +// +// created_at is stamped by the database. It is the order in which competing +// replicas take work, so it is measured on the one clock they share; a time taken +// in Go would make the order depend on which replica enqueued. +func EnqueueClaim(ctx context.Context, db *gorm.DB, kind messaging.WorkKind, payload any) (string, error) { + if err := requirePostgres(db, "enqueueing work"); err != nil { + return "", err + } + if !knownKind(kind) { + return "", fmt.Errorf("enqueueing work: unknown work kind %q", kind) + } + raw, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("encoding the payload of a %s claim: %w", kind, err) + } + if err := messaging.CheckWorkSize(kind, len(raw)); err != nil { + return "", err + } + id := uuid.New().String() + if err := db.WithContext(ctx).Model(&WorkClaim{}).Create(map[string]any{ + "id": id, + "kind": string(kind), + "state": ClaimPending, + "payload": raw, + "claimed_by": "", + "attempts": 0, + "created_at": gorm.Expr("now()"), + }).Error; err != nil { + return "", fmt.Errorf("enqueueing a %s claim: %w", kind, err) + } + return id, nil +} + +// claimNextSQL takes at most one pending row of a kind in one statement. +// +// FOR UPDATE SKIP LOCKED is the whole mechanism. Without SKIP LOCKED two +// replicas that poll together do not take two rows: the second blocks on the +// lock of the first and then looks again, so a slow claimant stalls every other +// replica. Without FOR UPDATE they can both read one row and both write it. +// It is one statement and not a select followed by an update for the same reason: +// the gap between the two is where two claimants see one row. +// +// claimed_at is stamped by the database, like created_at. +const claimNextSQL = `UPDATE ` + claimsTable + ` +SET state = '` + ClaimClaimed + `', claimed_by = ?, claimed_at = now() +WHERE id = ( + SELECT id FROM ` + claimsTable + ` + WHERE state = '` + ClaimPending + `' AND kind = ? + AND (not_before IS NULL OR not_before <= now()) + ORDER BY created_at, id + FOR UPDATE SKIP LOCKED + LIMIT 1 +) +RETURNING id, kind, state, payload, claimed_by, claimed_at, attempts, not_before, created_at` + +// ClaimNext takes the oldest pending row of kind and marks it as held by owner. +// It returns ErrNoWork when there is none. +func ClaimNext(ctx context.Context, db *gorm.DB, owner string, kind messaging.WorkKind) (*WorkClaim, error) { + if err := requirePostgres(db, "claiming work"); err != nil { + return nil, err + } + if owner == "" { + // A claim nobody owns cannot be reaped when its holder dies, because the + // reap asks whether the owner is live. The symptom of such a claim is + // work that is never retried and never answered. + return nil, errors.New("claiming work: no owner named, so the claim could never be reaped") + } + var claim WorkClaim + res := db.WithContext(ctx).Raw(claimNextSQL, owner, string(kind)).Scan(&claim) + if res.Error != nil { + return nil, fmt.Errorf("claiming work as %q: %w", owner, res.Error) + } + if res.RowsAffected == 0 { + return nil, ErrNoWork + } + return &claim, nil +} + +// The wait before a row that got no answer may be claimed again. +// +// There is no dead letter, and that is a decision. The only outcome that reaches +// ReleaseClaim is one where nothing was learned about the work: no agent worker +// was connected, the tunnel broke, a peer could not be reached. None of them is +// the worker saying that it ran the job and the job failed, and failing a job on +// any of them would report a missing connection as the verdict of a worker. A +// fleet that is down for a day must run its queued work when it comes back, and +// not find it failed. +// +// What is a defect is retrying at the poll interval for ever. At two seconds a +// row that can never be dispatched costs 43000 updates a day. Rows go oldest +// first, so it would also be claimed again ahead of every newer row on every +// tick and hold a slot while it fails. So the retry has no end and its rate has +// one: each release stamps the earliest time of the next claim, doubling from +// claimBackoffBase to claimBackoffCap. +const ( + // claimBackoffBase is the wait after the first release. + claimBackoffBase = 2 * time.Second + // claimBackoffCap bounds the wait, and with it how long queued work waits + // after a fleet comes back. + claimBackoffCap = 60 * time.Second + // claimBackoffMaxShift clamps the exponent, so that power() stays away from + // infinity for a row that has been retried for months. + claimBackoffMaxShift = 16 +) + +// releaseClaimSQL returns a row to the pool and stamps its next eligibility, in +// one statement and from the clock of the database. The row is matched on its +// holder and its attempt count, so a holder that was reaped cannot release the +// claim of the replica that took over. +const releaseClaimSQL = `UPDATE ` + claimsTable + ` +SET state = '` + ClaimPending + `', claimed_by = '', claimed_at = NULL, attempts = attempts + 1, + not_before = now() + make_interval(secs => LEAST(?, ? * power(2, LEAST(attempts, ?)))) +WHERE id = ? AND state = '` + ClaimClaimed + `' AND claimed_by = ? AND attempts = ?` + +// ReleaseClaim returns a claim to the pool with one more attempt and a wait. It +// is what a failure does that says nothing about the work. It reports false when +// the claim was no longer the holder's, which is not an error: another replica +// has the work now. +func ReleaseClaim(ctx context.Context, db *gorm.DB, claim *WorkClaim) (bool, error) { + if err := requirePostgres(db, "releasing a claim"); err != nil { + return false, err + } + res := db.WithContext(ctx).Exec(releaseClaimSQL, + claimBackoffCap.Seconds(), claimBackoffBase.Seconds(), claimBackoffMaxShift, + claim.ID, claim.ClaimedBy, claim.Attempts) + if res.Error != nil { + return false, fmt.Errorf("releasing claim %q: %w", claim.ID, res.Error) + } + return res.RowsAffected == 1, nil +} + +// CompleteClaim deletes the row. It is what an answer of the worker does, +// success or failure: the work ran and said what happened, so it must not run +// again. It reports false when the claim was no longer the holder's. +func CompleteClaim(ctx context.Context, db *gorm.DB, claim *WorkClaim) (bool, error) { + if db == nil { + return false, errors.New("completing a claim: no database handle") + } + res := db.WithContext(ctx). + Where("id = ? AND state = ? AND claimed_by = ? AND attempts = ?", claim.ID, ClaimClaimed, claim.ClaimedBy, claim.Attempts). + Delete(&WorkClaim{}) + if res.Error != nil { + return false, fmt.Errorf("completing claim %q: %w", claim.ID, res.Error) + } + return res.RowsAffected == 1, nil +} + +// reapAbandonedSQL releases every claim whose holder is not a live replica. +// +// The test is the absence of the replica and not the age of the claim. A claim +// that a heartbeating replica has held for an hour is a slow job, and taking it +// would run the work twice. A claim that a dead replica held for a second is work +// nobody drives. Age cannot tell the two apart, so there is no age in the +// statement: no time after which a live replica loses its work. +// +// It reads cluster.LiveInstanceIDsSQL and does not restate what live means, so +// this reap and every other reader of the absence of a replica move together. It +// is one statement, so a replica cannot die or return between the decision on who +// is live and the act. It stamps no wait: work whose replica died was never +// dispatched anywhere and has nothing to back off from. +const reapAbandonedSQL = `UPDATE ` + claimsTable + ` +SET state = '` + ClaimPending + `', claimed_by = '', claimed_at = NULL, attempts = attempts + 1, not_before = NULL +WHERE state = '` + ClaimClaimed + `' + AND claimed_by NOT IN (` + cluster.LiveInstanceIDsSQL + `)` + +// ReapAbandoned returns the claims of replicas that are no longer live to the +// pool, and says how many. liveness is the window of cluster.InstanceLiveness. It +// is not a timeout on the work. +func ReapAbandoned(ctx context.Context, db *gorm.DB, liveness time.Duration) (int64, error) { + if err := requirePostgres(db, "reaping abandoned claims"); err != nil { + return 0, err + } + res := db.WithContext(ctx).Exec(reapAbandonedSQL, liveness.Seconds()) + if res.Error != nil { + return 0, fmt.Errorf("reaping abandoned claims: %w", res.Error) + } + return res.RowsAffected, nil +} + +// ownerIsLiveSQL asks whether one replica is live, with the predicate of the +// reap. A replica that the reap would not keep must not claim: its claims would +// look like those of a dead replica, and another replica would take the work +// while it still runs. +const ownerIsLiveSQL = `SELECT count(*) FROM (` + cluster.LiveInstanceIDsSQL + ` AND instances.id = ?) AS live` + +// OwnerIsLive reports whether owner is a replica that the deployment considers +// live. A replica with no row in the instances table answers false, which is +// the honest answer: nothing about it is visible to its peers. +func OwnerIsLive(ctx context.Context, db *gorm.DB, owner string, liveness time.Duration) (bool, error) { + if err := requirePostgres(db, "checking whether this replica is live"); err != nil { + return false, err + } + var n int64 + if err := db.WithContext(ctx).Raw(ownerIsLiveSQL, liveness.Seconds(), owner).Scan(&n).Error; err != nil { + return false, fmt.Errorf("checking whether %q is live: %w", owner, err) + } + return n > 0, nil +} + +// migratePendingSQL takes the pending rows out of the queue. Rows that a replica +// holds are not touched: the replica drives them to their end where they are. +const migratePendingSQL = `UPDATE ` + claimsTable + ` +SET state = '` + ClaimMigrated + `' +WHERE state = '` + ClaimPending + `' +RETURNING id, kind, state, payload, claimed_by, claimed_at, attempts, not_before, created_at` + +// MigratePending takes every pending row out of the queue and returns it, for a +// change to a carrier that has its own queue. A row is migrated once: the +// statement is a compare-and-set on the state, so it never goes to a claimant and +// to the sweep, and a second sweep returns nothing. +// +// The caller publishes the rows after this returns. A sweeper that dies between +// the two loses the rows, as the queue of the new carrier loses them, and the +// alternative, to publish first, would run a job twice. +func MigratePending(ctx context.Context, db *gorm.DB) ([]*WorkClaim, error) { + if err := requirePostgres(db, "migrating pending work"); err != nil { + return nil, err + } + var rows []*WorkClaim + if err := db.WithContext(ctx).Raw(migratePendingSQL).Scan(&rows).Error; err != nil { + return nil, fmt.Errorf("migrating pending work: %w", err) + } + return rows, nil +} diff --git a/core/services/jobs/claim_test.go b/core/services/jobs/claim_test.go new file mode 100644 index 000000000000..b315af257f2e --- /dev/null +++ b/core/services/jobs/claim_test.go @@ -0,0 +1,493 @@ +package jobs + +import ( + "context" + "errors" + "path/filepath" + "runtime" + "strings" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + sqlitedriver "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// newClaimDB returns a database with the claim table and the cluster tables. +func newClaimDB() *gorm.DB { + GinkgoHelper() + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db := testutil.SetupTestDB() + Expect(cluster.Migrate(context.Background(), db)).To(Succeed()) + Expect(MigrateClaims(context.Background(), db)).To(Succeed()) + return db +} + +// live registers a replica, so that its claims are not reaped. +func live(db *gorm.DB, id string) { + GinkgoHelper() + Expect(cluster.NewRegistry(db).Register(context.Background(), id, "test", 0, "")).To(Succeed()) +} + +// kill makes a replica look as if it stopped heartbeating an hour ago. +func kill(db *gorm.DB, id string) { + GinkgoHelper() + Expect(db.Exec(`UPDATE instances SET last_seen = now() - interval '1 hour' WHERE id = ?`, id).Error).To(Succeed()) +} + +func rowOf(db *gorm.DB, id string) (WorkClaim, bool) { + GinkgoHelper() + var rows []WorkClaim + Expect(db.Where("id = ?", id).Find(&rows).Error).To(Succeed()) + if len(rows) == 0 { + return WorkClaim{}, false + } + return rows[0], true +} + +var _ = Describe("The claim table", func() { + var ( + ctx context.Context + db *gorm.DB + ) + + BeforeEach(func() { + ctx = context.Background() + db = newClaimDB() + live(db, "replica-a") + live(db, "replica-b") + }) + + Describe("MigrateClaims", func() { + It("can run again and again, also at the same time", func() { + var wg sync.WaitGroup + errs := make(chan error, 8) + for range 8 { + wg.Add(1) + go func() { defer wg.Done(); errs <- MigrateClaims(ctx, db) }() + } + wg.Wait() + close(errs) + for err := range errs { + Expect(err).ToNot(HaveOccurred()) + } + }) + }) + + Describe("EnqueueClaim", func() { + It("writes a pending row with the payload as JSON and the time of the database", func() { + id, err := EnqueueClaim(ctx, db, messaging.WorkMCPCI, map[string]string{"job_id": "j1"}) + Expect(err).ToNot(HaveOccurred()) + row, ok := rowOf(db, id) + Expect(ok).To(BeTrue()) + Expect(row.Kind).To(Equal("mcp-ci")) + Expect(row.State).To(Equal(ClaimPending)) + Expect(row.ClaimedBy).To(BeEmpty()) + Expect(row.Attempts).To(BeZero()) + Expect(row.NotBefore).To(BeNil()) + Expect(string(row.Payload)).To(MatchJSON(`{"job_id":"j1"}`)) + Expect(time.Since(row.CreatedAt)).To(BeNumerically("<", time.Minute)) + }) + + It("refuses a kind that is not one of the work kinds, and writes nothing", func() { + _, err := EnqueueClaim(ctx, db, messaging.WorkKind("nonsense"), 1) + Expect(err).To(HaveOccurred()) + var n int64 + Expect(db.Model(&WorkClaim{}).Count(&n).Error).To(Succeed()) + Expect(n).To(BeZero()) + }) + + It("refuses a payload above the bound that every carrier keeps", func() { + _, err := EnqueueClaim(ctx, db, messaging.WorkAgentRun, strings.Repeat("a", messaging.MaxWorkPayloadBytes+1)) + Expect(errors.Is(err, messaging.ErrWorkPayloadTooLarge)).To(BeTrue(), "got %v", err) + }) + + It("refuses a payload that does not encode", func() { + _, err := EnqueueClaim(ctx, db, messaging.WorkAgentRun, make(chan int)) + Expect(err).To(HaveOccurred()) + }) + }) + + Describe("ClaimNext", func() { + It("takes the oldest row of the kind and marks it as held by the owner", func() { + first, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, "first") + _, _ = EnqueueClaim(ctx, db, messaging.WorkMCPCI, "second") + _, _ = EnqueueClaim(ctx, db, messaging.WorkAgentRun, "other kind") + + claim, err := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(err).ToNot(HaveOccurred()) + Expect(claim.ID).To(Equal(first)) + Expect(claim.State).To(Equal(ClaimClaimed)) + Expect(claim.ClaimedBy).To(Equal("replica-a")) + Expect(claim.ClaimedAt).ToNot(BeNil()) + Expect(string(claim.Payload)).To(Equal(`"first"`)) + }) + + It("serves a kind in the order it was enqueued", func() { + var want []string + for i := range 12 { + id, err := EnqueueClaim(ctx, db, messaging.WorkAgentRun, i) + Expect(err).ToNot(HaveOccurred()) + want = append(want, id) + } + var got []string + for range 12 { + claim, err := ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + Expect(err).ToNot(HaveOccurred()) + got = append(got, claim.ID) + } + Expect(got).To(Equal(want)) + }) + + It("reports that there is no work, as an ordinary answer", func() { + _, err := ClaimNext(ctx, db, "replica-a", messaging.WorkTask) + Expect(err).To(MatchError(ErrNoWork)) + _, _ = EnqueueClaim(ctx, db, messaging.WorkMCPCI, 1) + _, err = ClaimNext(ctx, db, "replica-a", messaging.WorkTask) + Expect(err).To(MatchError(ErrNoWork)) + }) + + It("refuses an owner with no name, because a claim nobody owns could never be reaped", func() { + _, _ = EnqueueClaim(ctx, db, messaging.WorkTask, 1) + _, err := ClaimNext(ctx, db, "", messaging.WorkTask) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrNoWork)).To(BeFalse()) + }) + + It("gives a row to one of several claimants that ask at the same time", func() { + const rows, claimants = 24, 12 + for i := range rows { + _, err := EnqueueClaim(ctx, db, messaging.WorkMCPCI, i) + Expect(err).ToNot(HaveOccurred()) + } + var mu sync.Mutex + taken := map[string]int{} + var wg sync.WaitGroup + for c := range claimants { + wg.Add(1) + go func() { + defer GinkgoRecover() + defer wg.Done() + owner := "replica-a" + if c%2 == 1 { + owner = "replica-b" + } + for { + claim, err := ClaimNext(ctx, db, owner, messaging.WorkMCPCI) + if errors.Is(err, ErrNoWork) { + return + } + Expect(err).ToNot(HaveOccurred()) + mu.Lock() + taken[claim.ID]++ + mu.Unlock() + } + }() + } + wg.Wait() + Expect(taken).To(HaveLen(rows)) + for id, n := range taken { + Expect(n).To(Equal(1), id) + } + }) + }) + + Describe("CompleteClaim", func() { + It("deletes the row", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + claim, _ := ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + ok, err := CompleteClaim(ctx, db, claim) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + _, there := rowOf(db, id) + Expect(there).To(BeFalse()) + }) + + It("does nothing for a claim that another replica holds now", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + stale, _ := ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + kill(db, "replica-a") + _, err := ReapAbandoned(ctx, db, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + fresh, err := ClaimNext(ctx, db, "replica-b", messaging.WorkAgentRun) + Expect(err).ToNot(HaveOccurred()) + Expect(fresh.ID).To(Equal(id)) + + ok, err := CompleteClaim(ctx, db, stale) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeFalse(), "a late answer of the old holder must not delete the claim of the new one") + row, there := rowOf(db, id) + Expect(there).To(BeTrue()) + Expect(row.ClaimedBy).To(Equal("replica-b")) + }) + + It("does nothing for a claim that the same replica holds again after a reap", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + stale, _ := ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + kill(db, "replica-a") + _, _ = ReapAbandoned(ctx, db, cluster.InstanceLiveness) + live(db, "replica-a") + again, err := ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + Expect(err).ToNot(HaveOccurred()) + Expect(again.Attempts).To(Equal(stale.Attempts + 1)) + + ok, _ := CompleteClaim(ctx, db, stale) + Expect(ok).To(BeFalse()) + _, there := rowOf(db, id) + Expect(there).To(BeTrue()) + }) + }) + + Describe("ReleaseClaim", func() { + It("returns the row to the pool with one more attempt and a time before which it is not claimable", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, 1) + claim, _ := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + ok, err := ReleaseClaim(ctx, db, claim) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + + row, _ := rowOf(db, id) + Expect(row.State).To(Equal(ClaimPending)) + Expect(row.ClaimedBy).To(BeEmpty()) + Expect(row.ClaimedAt).To(BeNil()) + Expect(row.Attempts).To(Equal(1)) + Expect(row.NotBefore).ToNot(BeNil()) + + _, err = ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(err).To(MatchError(ErrNoWork), "the backoff keeps a row that just failed from coming straight back") + }) + + It("doubles the wait for each attempt up to the cap, from the clock of the database", func() { + _, _ = EnqueueClaim(ctx, db, messaging.WorkMCPCI, 1) + var waits []time.Duration + for range 8 { + Expect(db.Exec(`UPDATE work_claims SET not_before = NULL`).Error).To(Succeed()) + claim, err := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(err).ToNot(HaveOccurred()) + _, err = ReleaseClaim(ctx, db, claim) + Expect(err).ToNot(HaveOccurred()) + var secs float64 + Expect(db.Raw(`SELECT extract(epoch FROM not_before - now()) FROM work_claims`).Scan(&secs).Error).To(Succeed()) + waits = append(waits, time.Duration(secs*float64(time.Second))) + } + Expect(waits[0]).To(BeNumerically("~", claimBackoffBase, time.Second)) + Expect(waits[1]).To(BeNumerically("~", 2*claimBackoffBase, time.Second)) + Expect(waits[2]).To(BeNumerically("~", 4*claimBackoffBase, time.Second)) + for _, w := range waits { + Expect(w).To(BeNumerically("<=", claimBackoffCap+time.Second)) + } + Expect(waits[len(waits)-1]).To(BeNumerically("~", claimBackoffCap, time.Second)) + }) + + It("does nothing for a claim that another replica holds now", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, 1) + stale, _ := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + kill(db, "replica-a") + _, _ = ReapAbandoned(ctx, db, cluster.InstanceLiveness) + _, err := ClaimNext(ctx, db, "replica-b", messaging.WorkMCPCI) + Expect(err).ToNot(HaveOccurred()) + + ok, err := ReleaseClaim(ctx, db, stale) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeFalse()) + row, _ := rowOf(db, id) + Expect(row.State).To(Equal(ClaimClaimed)) + Expect(row.ClaimedBy).To(Equal("replica-b")) + }) + + It("lets a poison row wait while the rows behind it are served", func() { + poison, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, "poison") + var rest []string + for i := range 5 { + id, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, i) + rest = append(rest, id) + } + claim, _ := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(claim.ID).To(Equal(poison)) + _, _ = ReleaseClaim(ctx, db, claim) + + var got []string + for range 5 { + c, err := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(err).ToNot(HaveOccurred()) + got = append(got, c.ID) + } + Expect(got).To(Equal(rest)) + _, err := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(err).To(MatchError(ErrNoWork)) + }) + }) + + Describe("ReapAbandoned", func() { + It("releases the claims of a replica that is no longer live, with no wait", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + _, _ = ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + kill(db, "replica-a") + + n, err := ReapAbandoned(ctx, db, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeEquivalentTo(1)) + row, _ := rowOf(db, id) + Expect(row.State).To(Equal(ClaimPending)) + Expect(row.Attempts).To(Equal(1)) + Expect(row.NotBefore).To(BeNil(), "work that was never dispatched anywhere has nothing to back off from") + + got, err := ClaimNext(ctx, db, "replica-b", messaging.WorkAgentRun) + Expect(err).ToNot(HaveOccurred()) + Expect(got.ID).To(Equal(id)) + }) + + It("leaves the claims of a replica that is live, however long they have been held", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + _, _ = ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + Expect(db.Exec(`UPDATE work_claims SET claimed_at = now() - interval '1 day'`).Error).To(Succeed()) + + n, err := ReapAbandoned(ctx, db, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeZero()) + row, _ := rowOf(db, id) + Expect(row.State).To(Equal(ClaimClaimed)) + }) + + It("treats a replica that has no row at all as gone", func() { + id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + _, _ = ClaimNext(ctx, db, "replica-a", messaging.WorkAgentRun) + Expect(db.Exec(`DELETE FROM instances WHERE id = 'replica-a'`).Error).To(Succeed()) + n, err := ReapAbandoned(ctx, db, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeEquivalentTo(1)) + row, _ := rowOf(db, id) + Expect(row.State).To(Equal(ClaimPending)) + }) + + It("leaves pending and migrated rows alone", func() { + _, _ = EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) + n, err := ReapAbandoned(ctx, db, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(BeZero()) + }) + }) + + Describe("OwnerIsLive", func() { + It("is true for a registered replica and false for one that stopped or never registered", func() { + ok, err := OwnerIsLive(ctx, db, "replica-a", cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + kill(db, "replica-a") + ok, _ = OwnerIsLive(ctx, db, "replica-a", cluster.InstanceLiveness) + Expect(ok).To(BeFalse()) + ok, _ = OwnerIsLive(ctx, db, "stranger", cluster.InstanceLiveness) + Expect(ok).To(BeFalse()) + }) + }) + + Describe("MigratePending", func() { + It("takes every pending row out of the queue and returns it, and leaves the held rows", func() { + a, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, "a") + b, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, "b") + held, _ := EnqueueClaim(ctx, db, messaging.WorkMCPCI, "held") + // The oldest MCP CI row is "a"; take it and put the third one back in + // the order of the table so that "held" is the one that stays held. + Expect(db.Exec(`UPDATE work_claims SET created_at = now() + interval '1 minute' WHERE id = ?`, held).Error).To(Succeed()) + claim, err := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + Expect(err).ToNot(HaveOccurred()) + Expect(claim.ID).To(Equal(a)) + + moved, err := MigratePending(ctx, db) + Expect(err).ToNot(HaveOccurred()) + var ids []string + for _, c := range moved { + ids = append(ids, c.ID) + } + Expect(ids).To(ConsistOf(b, held)) + for _, c := range moved { + Expect(c.Kind).ToNot(BeEmpty()) + Expect(c.Payload).ToNot(BeEmpty()) + } + + for _, kind := range []messaging.WorkKind{messaging.WorkMCPCI, messaging.WorkAgentRun} { + _, err := ClaimNext(ctx, db, "replica-b", kind) + Expect(err).To(MatchError(ErrNoWork), "a migrated row is not claimable: %s", kind) + } + row, _ := rowOf(db, a) + Expect(row.State).To(Equal(ClaimClaimed), "a row that a replica holds is driven to its end where it is") + row, _ = rowOf(db, b) + Expect(row.State).To(Equal(ClaimMigrated)) + }) + + It("returns nothing the second time, so a repeated sweep moves nothing twice", func() { + _, _ = EnqueueClaim(ctx, db, messaging.WorkMCPCI, 1) + first, err := MigratePending(ctx, db) + Expect(err).ToNot(HaveOccurred()) + Expect(first).To(HaveLen(1)) + second, err := MigratePending(ctx, db) + Expect(err).ToNot(HaveOccurred()) + Expect(second).To(BeEmpty()) + }) + + It("gives each row to the sweep or to a claimant and never to both", func() { + const rows = 40 + for i := range rows { + _, _ = EnqueueClaim(ctx, db, messaging.WorkMCPCI, i) + } + var mu sync.Mutex + owners := map[string]string{} + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer GinkgoRecover() + defer wg.Done() + for { + c, err := ClaimNext(ctx, db, "replica-a", messaging.WorkMCPCI) + if errors.Is(err, ErrNoWork) { + return + } + Expect(err).ToNot(HaveOccurred()) + mu.Lock() + Expect(owners).ToNot(HaveKey(c.ID)) + owners[c.ID] = "claimed" + mu.Unlock() + } + }() + go func() { + defer GinkgoRecover() + defer wg.Done() + for range 20 { + moved, err := MigratePending(ctx, db) + Expect(err).ToNot(HaveOccurred()) + mu.Lock() + for _, c := range moved { + Expect(owners).ToNot(HaveKey(c.ID)) + owners[c.ID] = "migrated" + } + mu.Unlock() + time.Sleep(time.Millisecond) + } + }() + wg.Wait() + moved, _ := MigratePending(ctx, db) + for _, c := range moved { + Expect(owners).ToNot(HaveKey(c.ID)) + owners[c.ID] = "migrated" + } + Expect(owners).To(HaveLen(rows)) + }) + }) + + It("refuses a database that is not PostgreSQL, and says why", func() { + sqlite, err := gorm.Open(sqlitedriver.Open(filepath.Join(GinkgoT().TempDir(), "claims.db")), &gorm.Config{}) + Expect(err).ToNot(HaveOccurred()) + _, err = EnqueueClaim(ctx, sqlite, messaging.WorkTask, 1) + Expect(err).To(MatchError(ContainSubstring("requires PostgreSQL"))) + _, err = ClaimNext(ctx, sqlite, "replica-a", messaging.WorkTask) + Expect(err).To(MatchError(ContainSubstring("requires PostgreSQL"))) + }) +}) From 95b07ee61dfb117fe8d3ce9a7338f692855e3b7f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:15:41 +0000 Subject: [PATCH 080/162] feat(jobs): add the claim queue as a WorkQueue and a consumer for frontend replicas ClaimQueue puts a payload in the claim table and sends a wake hint on a broadcast subject. ClaimConsumer is the other half: it polls every two seconds, wakes at once on a hint for its kind, and claims into as many slots as it was given. The hint is at-most-once and may be lost or refused, and the poll finds the work then. With a poll of two seconds, the median wait of a unit is under 50 ms because of the hint. The outcome of the handler settles the claim: nil completes it, an error or a panic releases it with a wait, and an error that wraps ErrKeepClaim leaves it with the replica, for work that ran and whose answer could not be recorded, since releasing it would repeat the work. A replica whose own row is not live claims nothing, because its claims would look abandoned. Each poll first returns the claims of departed replicas to the queue. The queue passes the same conformance suite as the queue groups of NATS. The specs also cover what the suite leaves to the carrier: a replica that dies in the middle of a job loses nothing, a slow job of a live replica is not taken, a poison unit waits while the units behind it are served, and the order of a kind holds. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/jobs/claim_queue.go | 346 +++++++++++++++ core/services/jobs/claim_queue_test.go | 585 +++++++++++++++++++++++++ core/services/messaging/subjects.go | 6 + 3 files changed, 937 insertions(+) create mode 100644 core/services/jobs/claim_queue.go create mode 100644 core/services/jobs/claim_queue_test.go diff --git a/core/services/jobs/claim_queue.go b/core/services/jobs/claim_queue.go new file mode 100644 index 000000000000..95e871ea1791 --- /dev/null +++ b/core/services/jobs/claim_queue.go @@ -0,0 +1,346 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "time" + + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/xlog" +) + +// ClaimWake is the payload of a wake hint. It names the kind, so that the +// consumers of the other kinds do not look. +type ClaimWake struct { + Kind messaging.WorkKind `json:"kind"` +} + +// ClaimQueue is the WorkQueue of the carrier that has no broker: a payload +// becomes a row, and a consumer takes the row. +// +// Enqueue succeeds while nobody consumes, and the row waits. That differs from a +// publish to an empty queue group, which is lost. +type ClaimQueue struct { + db *gorm.DB + hints messaging.Publisher +} + +// NewClaimQueue returns a queue over db. hints, when not nil, gets a wake hint +// for each unit put in the queue. It may fail or lose the hint: the poll of the +// consumers is what guarantees that the work is found. +func NewClaimQueue(db *gorm.DB, hints messaging.Publisher) *ClaimQueue { + return &ClaimQueue{db: db, hints: hints} +} + +var _ messaging.WorkQueue = (*ClaimQueue)(nil) + +// Enqueue writes the payload as a row, and then tells the consumers. A refusal of +// the payload (its kind, its size, its encoding) comes back as an error and +// writes nothing. A hint that cannot be published is not an error. +func (q *ClaimQueue) Enqueue(ctx context.Context, kind messaging.WorkKind, payload any) error { + if _, err := EnqueueClaim(ctx, q.db, kind, payload); err != nil { + return err + } + if q.hints != nil { + if err := q.hints.Publish(messaging.SubjectClaimWake, ClaimWake{Kind: kind}); err != nil { + xlog.Debug("A wake hint for the claim queue could not be published; the poll will find the work", "kind", kind, "error", err) + } + } + return nil +} + +// ErrKeepClaim, wrapped in the error of a handler, tells the consumer to leave +// the claim as it is: neither delete it nor release it. It is for work that ran +// and whose answer could not be recorded. Releasing the claim would run the work +// again. The row stays held by this replica, and it is taken over only if the +// replica dies. +var ErrKeepClaim = errors.New("jobs: keep the claim") + +// settleTimeout bounds the writes that settle a claim. They run on a context that +// the stop of the replica does not cancel, because a replica that is leaving must +// still release what it holds. +const settleTimeout = 10 * time.Second + +// ClaimConsumerConfig is what a ClaimConsumer needs. +type ClaimConsumerConfig struct { + DB *gorm.DB + // Owner is the id of this replica in the instances table. The reap asks + // whether the owner of a claim is live, so the replica must be registered. + Owner string + // Events is what the handlers get as the place for their events. + Events messaging.Publisher + // Hints is where the wake hints are heard. It may be nil: the consumer then + // polls only. + Hints messaging.Broadcaster + // Interval is the poll. DefaultClaimInterval when zero. + Interval time.Duration + // Liveness is the window of replica liveness. cluster.InstanceLiveness when + // zero: it is the one definition of absence that the cluster has. + Liveness time.Duration +} + +// DefaultClaimInterval is how often a consumer looks for work when no hint woke +// it. +const DefaultClaimInterval = 2 * time.Second + +// ClaimConsumer is the WorkConsumer of the claim queue, for a frontend replica. +// It claims the rows of a kind and runs a handler on each. On the tunnel carrier +// the handler hands the work to an agent worker, because an agent worker has no +// database and cannot claim. +// +// The outcome of a handler settles the row. A nil return completes it: the work +// has an answer. An error releases it with a wait: nothing was learned about the +// work, so it must be tried again, possibly by another replica. A panic is an +// error. An error that wraps ErrKeepClaim leaves the row alone. +// +// A handler can be called again for the same work: after the replica that held +// it died, after a release. It must tolerate a repeat. +type ClaimConsumer struct { + cfg ClaimConsumerConfig +} + +var _ messaging.WorkConsumer = (*ClaimConsumer)(nil) + +// NewClaimConsumer returns a consumer. It refuses a configuration that would +// claim work and never answer it. +func NewClaimConsumer(cfg ClaimConsumerConfig) (*ClaimConsumer, error) { + switch { + case cfg.DB == nil: + return nil, errors.New("the claim consumer was built with no database to claim work from") + case cfg.Owner == "": + return nil, errors.New("the claim consumer was built with no instance id: its claims could not be told from those of a dead replica") + case cfg.Events == nil: + return nil, errors.New("the claim consumer was built with no publisher for the events of its handlers") + } + if cfg.Interval <= 0 { + cfg.Interval = DefaultClaimInterval + } + if cfg.Liveness <= 0 { + cfg.Liveness = cluster.InstanceLiveness + } + return &ClaimConsumer{cfg: cfg}, nil +} + +// Consume claims the rows of kind and runs h on each, at most maxInFlight at a +// time (0 and negative values are unbounded). ctx is the parent of every handler +// call. Unsubscribe stops claiming and waits for the handlers that run, so a +// handler must not call it. +func (c *ClaimConsumer) Consume(ctx context.Context, kind messaging.WorkKind, maxInFlight int, h messaging.WorkHandler) (messaging.Subscription, error) { + if !knownKind(kind) { + return nil, fmt.Errorf("consuming work: unknown work kind %q", kind) + } + loopCtx, stop := context.WithCancel(ctx) + l := &claimLoop{ + cfg: c.cfg, + kind: kind, + handler: h, + parent: ctx, + stop: stop, + wake: make(chan struct{}, 1), + done: make(chan struct{}), + } + if maxInFlight > 0 { + l.slots = make(chan struct{}, maxInFlight) + } + if c.cfg.Hints != nil { + sub, err := c.cfg.Hints.Subscribe(messaging.SubjectClaimWake, l.hear) + if err != nil { + stop() + return nil, fmt.Errorf("subscribing to the wake hints of the claim queue: %w", err) + } + l.hints = sub + } + go l.run(loopCtx) + return l, nil +} + +// claimLoop is one subscription: the loop that claims the rows of one kind. +type claimLoop struct { + cfg ClaimConsumerConfig + kind messaging.WorkKind + handler messaging.WorkHandler + parent context.Context + stop context.CancelFunc + slots chan struct{} // nil: unbounded + hints messaging.Subscription + + wake chan struct{} + done chan struct{} + wg sync.WaitGroup + + closeOnce sync.Once + + // unregistered remembers that this replica has said that it cannot claim, so + // the refusal is one line when it begins and one when it ends, and not one + // line per poll for as long as the misconfiguration lasts. + unregistered bool +} + +// hear takes a wake hint. A hint for another kind is ignored; one that cannot be +// read is taken as for this kind, because a needless look costs one query. +func (l *claimLoop) hear(raw []byte) { + var w ClaimWake + if err := json.Unmarshal(raw, &w); err == nil && w.Kind != l.kind { + return + } + select { + case l.wake <- struct{}{}: + default: // a look is already asked for + } +} + +func (l *claimLoop) run(ctx context.Context) { + defer close(l.done) + ticker := time.NewTicker(l.cfg.Interval) + defer ticker.Stop() + l.fill(ctx, true) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + l.fill(ctx, true) + case <-l.wake: + l.fill(ctx, false) + } + } +} + +// take reserves a slot for a handler, without waiting. +func (l *claimLoop) take() bool { + if l.slots == nil { + return true + } + select { + case l.slots <- struct{}{}: + return true + default: + return false + } +} + +func (l *claimLoop) free() { + if l.slots != nil { + <-l.slots + } +} + +// fill claims rows into the free slots. On a poll it first looks after the +// claims that departed replicas left. +func (l *claimLoop) fill(ctx context.Context, poll bool) { + // The liveness of this replica is checked before it touches the queue. A + // replica with no row in the instances table is invisible to its peers. That + // is a degraded state, but not one for holding a claim: the reap would take + // the claim out from under it while the work still ran, and the work would + // run twice. Refusing to claim heals by itself when membership registers. + live, err := OwnerIsLive(ctx, l.cfg.DB, l.cfg.Owner, l.cfg.Liveness) + if err != nil { + if ctx.Err() == nil { + xlog.Warn("Could not check whether this replica may claim work", "kind", l.kind, "error", err) + } + return + } + if !live { + if !l.unregistered { + l.unregistered = true + xlog.Error("This replica is not live in the cluster, so it does not claim queued work: its peers could not tell its claims from those of a dead replica", + "instance", l.cfg.Owner, "kind", l.kind) + } + return + } + if l.unregistered { + l.unregistered = false + xlog.Info("This replica is live in the cluster again and claims queued work", "instance", l.cfg.Owner, "kind", l.kind) + } + + if poll { + if released, err := ReapAbandoned(ctx, l.cfg.DB, l.cfg.Liveness); err != nil { + // Reported and not returned: a reap that could not run says nothing + // about the claim this tick is about to take. + if ctx.Err() == nil { + xlog.Warn("Could not release the claims of departed replicas", "error", err) + } + } else if released > 0 { + xlog.Info("Released work claimed by replicas that are no longer live", "count", released) + } + } + + for ctx.Err() == nil { + if !l.take() { + return // this replica drives as much as it may + } + claim, err := ClaimNext(ctx, l.cfg.DB, l.cfg.Owner, l.kind) + if err != nil { + l.free() + if !errors.Is(err, ErrNoWork) && ctx.Err() == nil { + xlog.Warn("Could not claim work", "kind", l.kind, "error", err) + } + return + } + l.wg.Add(1) + go l.drive(claim) + } +} + +// drive runs the handler on a claimed row and settles the row. +func (l *claimLoop) drive(claim *WorkClaim) { + defer l.wg.Done() + defer l.free() + err := l.call(claim) + + // The settle must outlive the stop of the replica. + ctx, cancel := context.WithTimeout(context.WithoutCancel(l.parent), settleTimeout) + defer cancel() + switch { + case err == nil: + l.settle(CompleteClaim(ctx, l.cfg.DB, claim)) + case errors.Is(err, ErrKeepClaim): + xlog.Warn("Leaving a claim as it is: the work ran and its answer could not be recorded", + "claim", claim.ID, "kind", claim.Kind, "error", err) + default: + xlog.Warn("Releasing a claim whose work obtained no answer; it will be tried again after a wait", + "claim", claim.ID, "kind", claim.Kind, "attempt", claim.Attempts+1, "error", err) + l.settle(ReleaseClaim(ctx, l.cfg.DB, claim)) + } +} + +func (l *claimLoop) settle(taken bool, err error) { + switch { + case err != nil: + // The row stays claimed by this replica, which is the safe side: it is + // taken over if this replica dies. + xlog.Error("Could not settle a claim", "error", err) + case !taken: + xlog.Warn("A claim was no longer held by this replica when its work ended; another replica has taken it over") + } +} + +// call runs the handler and turns a panic into an error, so that one bad payload +// cannot take the replica down with it. +func (l *claimLoop) call(claim *WorkClaim) (err error) { + defer func() { + if p := recover(); p != nil { + err = fmt.Errorf("the handler panicked: %v", p) + } + }() + return l.handler(l.parent, claim.Payload, l.cfg.Events) +} + +// Unsubscribe stops claiming and waits for the handlers that run. +func (l *claimLoop) Unsubscribe() error { + l.closeOnce.Do(func() { + l.stop() + <-l.done + if l.hints != nil { + _ = l.hints.Unsubscribe() + } + l.wg.Wait() + }) + return nil +} diff --git a/core/services/jobs/claim_queue_test.go b/core/services/jobs/claim_queue_test.go new file mode 100644 index 000000000000..5b1a22376eab --- /dev/null +++ b/core/services/jobs/claim_queue_test.go @@ -0,0 +1,585 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/messaging/messagingtest" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// droppingBus accepts every publish and delivers none, as a bus that loses its +// messages does. +type droppingBus struct{ messaging.Broadcaster } + +func (droppingBus) Publish(string, any) error { return nil } + +// brokenBus refuses every publish. +type brokenBus struct{ messaging.Broadcaster } + +func (brokenBus) Publish(string, any) error { return errors.New("the bus is down") } + +type consumerOptions struct { + owner string + events messaging.Publisher + hints messaging.Broadcaster + interval time.Duration +} + +// newConsumer returns a consumer for a replica that is live. +func newConsumer(db *gorm.DB, o consumerOptions) *ClaimConsumer { + GinkgoHelper() + live(db, o.owner) + if o.events == nil { + o.events = testutil.NewFakeBus() + } + c, err := NewClaimConsumer(ClaimConsumerConfig{ + DB: db, Owner: o.owner, Events: o.events, Hints: o.hints, Interval: o.interval, + }) + Expect(err).ToNot(HaveOccurred()) + return c +} + +var _ = Describe("The claim queue as a WorkQueue", func() { + messagingtest.RunWorkQueueConformance(func() messagingtest.WorkQueueRig { + db := newClaimDB() + bus := testutil.NewFakeBus() + var n atomic.Int32 + return messagingtest.WorkQueueRig{ + Queue: NewClaimQueue(db, bus), + NewConsumer: func() messaging.WorkConsumer { + return newConsumer(db, consumerOptions{ + owner: fmt.Sprintf("replica-%d", n.Add(1)), + hints: bus, + interval: 100 * time.Millisecond, + }) + }, + } + }) +}) + +var _ = Describe("The claim queue", func() { + var ( + ctx context.Context + db *gorm.DB + bus *testutil.FakeBus + ) + + BeforeEach(func() { + c, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + ctx = c + db = newClaimDB() + bus = testutil.NewFakeBus() + }) + + consume := func(c *ClaimConsumer, kind messaging.WorkKind, max int, h messaging.WorkHandler) messaging.Subscription { + GinkgoHelper() + sub, err := c.Consume(ctx, kind, max, h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + return sub + } + + rows := func() int64 { + var n int64 + Expect(db.Model(&WorkClaim{}).Count(&n).Error).To(Succeed()) + return n + } + + Describe("the wake hint", func() { + It("starts a consumer at once, with a poll interval that would have kept it waiting for seconds", func() { + picked := make(chan time.Time, 64) + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 2 * time.Second}) + consume(c, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + picked <- time.Now() + return nil + }) + queue := NewClaimQueue(db, bus) + + var waits []time.Duration + for i := range 20 { + time.Sleep(20 * time.Millisecond) + start := time.Now() + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, i)).To(Succeed()) + var at time.Time + Eventually(picked, 5*time.Second).Should(Receive(&at)) + waits = append(waits, at.Sub(start)) + } + sort.Slice(waits, func(i, j int) bool { return waits[i] < waits[j] }) + Expect(waits[len(waits)/2]).To(BeNumerically("<", 50*time.Millisecond), "the p50 of the pickup, waits: %v", waits) + }) + + It("loses no work when every hint is dropped, because the poll finds it", func() { + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 150 * time.Millisecond}) + consume(c, messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { + got.Add(1) + return nil + }) + queue := NewClaimQueue(db, droppingBus{bus}) + for i := range 10 { + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, i)).To(Succeed()) + } + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(10)) + Eventually(rows, 5*time.Second).Should(BeZero()) + }) + + It("enqueues when the bus is down, and says nothing of the hint", func() { + queue := NewClaimQueue(db, brokenBus{bus}) + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, "x")).To(Succeed()) + Expect(rows()).To(BeEquivalentTo(1)) + }) + + It("enqueues with no bus at all", func() { + Expect(NewClaimQueue(db, nil).Enqueue(ctx, messaging.WorkMCPCI, "x")).To(Succeed()) + Expect(rows()).To(BeEquivalentTo(1)) + }) + + It("names the kind in the hint, on the subject that every carrier serves", func() { + var got []byte + _, err := bus.Subscribe(messaging.SubjectClaimWake, func(b []byte) { got = b }) + Expect(err).ToNot(HaveOccurred()) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkAgentRun, 1)).To(Succeed()) + Expect(got).To(MatchJSON(`{"kind":"agent-run"}`)) + Expect(messaging.ValidateBroadcastSubject(messaging.SubjectClaimWake)).To(Succeed()) + }) + + It("does not wake the consumers of another kind", func() { + // The poll interval is long, so only a hint could start the handler. + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: time.Hour}) + consume(c, messaging.WorkTask, 0, func(context.Context, []byte, messaging.Publisher) error { + got.Add(1) + return nil + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "x")).To(Succeed()) + Consistently(got.Load, 300*time.Millisecond).Should(BeZero()) + }) + }) + + Describe("work that nobody consumes yet", func() { + It("stays in the table, and a consumer that starts later gets it", func() { + queue := NewClaimQueue(db, bus) + for i := range 5 { + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, i)).To(Succeed()) + } + Expect(rows()).To(BeEquivalentTo(5)) + + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(c, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + got.Add(1) + return nil + }) + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(5)) + }) + }) + + Describe("exactly once under competing consumers", func() { + It("runs every unit once when five replicas with four slots each compete for two hundred", func() { + const items = 200 + var mu sync.Mutex + seen := map[string]int{} + handler := func(_ context.Context, p []byte, _ messaging.Publisher) error { + var j struct{ N int } + Expect(json.Unmarshal(p, &j)).To(Succeed()) + time.Sleep(time.Millisecond) + mu.Lock() + seen[fmt.Sprint(j.N)]++ + mu.Unlock() + return nil + } + for r := range 5 { + c := newConsumer(db, consumerOptions{owner: fmt.Sprintf("replica-%d", r), hints: bus, interval: 100 * time.Millisecond}) + consume(c, messaging.WorkAgentRun, 4, handler) + } + queue := NewClaimQueue(db, bus) + for i := range items { + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, map[string]int{"N": i})).To(Succeed()) + } + Eventually(func() int { mu.Lock(); defer mu.Unlock(); return len(seen) }, 30*time.Second).Should(Equal(items)) + Eventually(rows, 10*time.Second).Should(BeZero()) + mu.Lock() + defer mu.Unlock() + for n, count := range seen { + Expect(count).To(Equal(1), "unit %s", n) + } + }) + }) + + Describe("order", func() { + It("hands the units of a kind to a serial consumer in the order they were enqueued", func() { + var mu sync.Mutex + var order []int + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + queue := NewClaimQueue(db, bus) + for i := range 25 { + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, i)).To(Succeed()) + } + consume(c, messaging.WorkMCPCI, 1, func(_ context.Context, p []byte, _ messaging.Publisher) error { + var n int + Expect(json.Unmarshal(p, &n)).To(Succeed()) + mu.Lock() + order = append(order, n) + mu.Unlock() + return nil + }) + Eventually(func() int { mu.Lock(); defer mu.Unlock(); return len(order) }, 15*time.Second).Should(Equal(25)) + mu.Lock() + defer mu.Unlock() + Expect(sort.IntsAreSorted(order)).To(BeTrue(), "got %v", order) + }) + + It("keeps the order of each kind apart from the others", func() { + queue := NewClaimQueue(db, bus) + for i := range 6 { + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, i)).To(Succeed()) + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, 100+i)).To(Succeed()) + } + var mu sync.Mutex + got := map[messaging.WorkKind][]int{} + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + for _, kind := range []messaging.WorkKind{messaging.WorkMCPCI, messaging.WorkAgentRun} { + consume(c, kind, 1, func(_ context.Context, p []byte, _ messaging.Publisher) error { + var n int + _ = json.Unmarshal(p, &n) + mu.Lock() + got[kind] = append(got[kind], n) + mu.Unlock() + return nil + }) + } + Eventually(func() int { mu.Lock(); defer mu.Unlock(); return len(got[messaging.WorkMCPCI]) + len(got[messaging.WorkAgentRun]) }, 15*time.Second).Should(Equal(12)) + mu.Lock() + defer mu.Unlock() + Expect(got[messaging.WorkMCPCI]).To(Equal([]int{0, 1, 2, 3, 4, 5})) + Expect(got[messaging.WorkAgentRun]).To(Equal([]int{100, 101, 102, 103, 104, 105})) + }) + }) + + Describe("a replica that dies in the middle of a job", func() { + It("returns the job to the queue and runs it on another replica, with no loss", func() { + started := make(chan struct{}) + hold := make(chan struct{}) + var ranOnA, ranOnB atomic.Int32 + + a := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(a, messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { + ranOnA.Add(1) + close(started) + <-hold // the process is gone: this never returns during the spec + return nil + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkAgentRun, "important")).To(Succeed()) + Eventually(started, 5*time.Second).Should(BeClosed()) + Expect(rows()).To(BeEquivalentTo(1)) + + // A stops heartbeating. B is another replica, and it was there all along. + kill(db, "replica-a") + b := newConsumer(db, consumerOptions{owner: "replica-b", hints: bus, interval: 100 * time.Millisecond}) + consume(b, messaging.WorkAgentRun, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + ranOnB.Add(1) + Expect(string(p)).To(Equal(`"important"`)) + return nil + }) + Eventually(ranOnB.Load, 10*time.Second).Should(BeEquivalentTo(1)) + Eventually(rows, 5*time.Second).Should(BeZero()) + + // The late answer of A changes nothing. + close(hold) + Consistently(rows, 300*time.Millisecond).Should(BeZero()) + Expect(ranOnA.Load()).To(BeEquivalentTo(1)) + Expect(ranOnB.Load()).To(BeEquivalentTo(1), "the job ran once more on B, which is what at-least-once means") + }) + + It("does not take a job from a replica that only runs it slowly", func() { + started := make(chan struct{}) + hold := make(chan struct{}) + var ranOnB atomic.Int32 + a := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(a, messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { + close(started) + <-hold + return nil + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkAgentRun, "slow")).To(Succeed()) + Eventually(started, 5*time.Second).Should(BeClosed()) + Expect(db.Exec(`UPDATE work_claims SET claimed_at = now() - interval '2 hours'`).Error).To(Succeed()) + + b := newConsumer(db, consumerOptions{owner: "replica-b", hints: bus, interval: 50 * time.Millisecond}) + consume(b, messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { + ranOnB.Add(1) + return nil + }) + Consistently(ranOnB.Load, time.Second).Should(BeZero()) + close(hold) + Eventually(rows, 5*time.Second).Should(BeZero()) + }) + + It("recovers every job it held, and the ones still waiting go to the others", func() { + const jobs = 12 + hold := make(chan struct{}) + defer close(hold) + var held atomic.Int32 + a := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(a, messaging.WorkMCPCI, 3, func(context.Context, []byte, messaging.Publisher) error { + held.Add(1) + <-hold + return nil + }) + queue := NewClaimQueue(db, bus) + for i := range jobs { + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, i)).To(Succeed()) + } + Eventually(held.Load, 5*time.Second).Should(BeEquivalentTo(3)) + kill(db, "replica-a") + + var mu sync.Mutex + ran := map[int]int{} + b := newConsumer(db, consumerOptions{owner: "replica-b", hints: bus, interval: 100 * time.Millisecond}) + consume(b, messaging.WorkMCPCI, 4, func(_ context.Context, p []byte, _ messaging.Publisher) error { + var n int + _ = json.Unmarshal(p, &n) + mu.Lock() + ran[n]++ + mu.Unlock() + return nil + }) + Eventually(func() int { mu.Lock(); defer mu.Unlock(); return len(ran) }, 15*time.Second).Should(Equal(jobs)) + Eventually(rows, 5*time.Second).Should(BeZero()) + }) + + It("claims nothing while its own row says that it is not live", func() { + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + kill(db, "replica-a") + consume(c, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + got.Add(1) + return nil + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, 1)).To(Succeed()) + Consistently(got.Load, 500*time.Millisecond).Should(BeZero()) + + live(db, "replica-a") // its membership registers + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(1)) + }) + }) + + Describe("a unit of work that fails", func() { + It("goes back to the queue with a wait, and the work behind it is not held up", func() { + var mu sync.Mutex + runs := map[string]int{} + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + consume(c, messaging.WorkMCPCI, 1, func(_ context.Context, p []byte, _ messaging.Publisher) error { + mu.Lock() + runs[string(p)]++ + mu.Unlock() + if string(p) == `"poison"` { + return errors.New("cannot serve") + } + return nil + }) + queue := NewClaimQueue(db, bus) + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, "poison")).To(Succeed()) + for i := range 5 { + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, fmt.Sprintf("ok-%d", i))).To(Succeed()) + } + Eventually(func() int { + mu.Lock() + defer mu.Unlock() + n := 0 + for k, v := range runs { + if k != `"poison"` { + n += v + } + } + return n + }, 10*time.Second).Should(Equal(5)) + + Consistently(func() int { mu.Lock(); defer mu.Unlock(); return runs[`"poison"`] }, time.Second).Should(BeNumerically("<=", 2), + "a wait of at least two seconds keeps it from coming back at the poll interval") + Expect(rows()).To(BeEquivalentTo(1)) + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed()) + Expect(row.Attempts).To(BeNumerically(">=", 1)) + Expect(row.State).To(Equal(ClaimPending)) + }) + + It("survives a handler that panics, and gives the row the same wait", func() { + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + consume(c, messaging.WorkAgentRun, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { + if string(p) == `"bomb"` { + panic("bug in the handler") + } + got.Add(1) + return nil + }) + queue := NewClaimQueue(db, bus) + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, "bomb")).To(Succeed()) + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, "fine")).To(Succeed()) + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(1)) + Eventually(func() int { + var row WorkClaim + if err := db.First(&row).Error; err != nil { + return -1 + } + return row.Attempts + }, 5*time.Second).Should(BeNumerically(">=", 1)) + }) + + It("is dropped for good when the handler says it could not read the payload, as a handler returns nil for that", func() { + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + consume(c, messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { + got.Add(1) + return nil + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkAgentRun, "garbage")).To(Succeed()) + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(1)) + Eventually(rows, 5*time.Second).Should(BeZero()) + Consistently(got.Load, 300*time.Millisecond).Should(BeEquivalentTo(1)) + }) + + It("stays with its replica when the answer could not be recorded, because running it again would repeat the work", func() { + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + consume(c, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + got.Add(1) + return fmt.Errorf("the store refused the result: %w", ErrKeepClaim) + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, 1)).To(Succeed()) + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(1)) + Consistently(got.Load, time.Second).Should(BeEquivalentTo(1)) + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed()) + Expect(row.State).To(Equal(ClaimClaimed)) + Expect(row.ClaimedBy).To(Equal("replica-a")) + + // If the replica dies, the work is taken over: the answer is lost with it. + kill(db, "replica-a") + var ranOnB atomic.Int32 + b := newConsumer(db, consumerOptions{owner: "replica-b", hints: bus, interval: 50 * time.Millisecond}) + consume(b, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + ranOnB.Add(1) + return nil + }) + Eventually(ranOnB.Load, 5*time.Second).Should(BeEquivalentTo(1)) + }) + }) + + Describe("a consumer", func() { + It("hands its handlers the publisher it was built with, and the context it was given", func() { + events := testutil.NewFakeBus() + type ctxKey struct{} + parent := context.WithValue(ctx, ctxKey{}, "marker") + var gotEvents messaging.Publisher + var marker atomic.Value + done := make(chan struct{}) + c := newConsumer(db, consumerOptions{owner: "replica-a", events: events, hints: bus, interval: 50 * time.Millisecond}) + sub, err := c.Consume(parent, messaging.WorkMCPCI, 0, func(ctx context.Context, _ []byte, e messaging.Publisher) error { + gotEvents = e + marker.Store(ctx.Value(ctxKey{})) + close(done) + return nil + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, 1)).To(Succeed()) + Eventually(done, 5*time.Second).Should(BeClosed()) + Expect(gotEvents).To(BeIdenticalTo(events)) + Expect(marker.Load()).To(Equal("marker")) + }) + + It("settles the claim of a handler that runs when the replica is told to stop, and Unsubscribe waits for it", func() { + started := make(chan struct{}) + release := make(chan struct{}) + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + sub, err := c.Consume(ctx, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + close(started) + <-release + return nil + }) + Expect(err).ToNot(HaveOccurred()) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, 1)).To(Succeed()) + Eventually(started, 5*time.Second).Should(BeClosed()) + + stopped := make(chan struct{}) + go func() { defer close(stopped); _ = sub.Unsubscribe() }() + Consistently(stopped, 200*time.Millisecond).ShouldNot(BeClosed()) + close(release) + Eventually(stopped, 5*time.Second).Should(BeClosed()) + Expect(rows()).To(BeZero(), "the claim was settled before Unsubscribe returned") + }) + + It("settles a claim after the context of the replica ended, so a restart does not leave it held", func() { + started := make(chan struct{}) + parent, stop := context.WithCancel(ctx) + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + sub, err := c.Consume(parent, messaging.WorkMCPCI, 0, func(ctx context.Context, _ []byte, _ messaging.Publisher) error { + close(started) + <-ctx.Done() + return ctx.Err() + }) + Expect(err).ToNot(HaveOccurred()) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, 1)).To(Succeed()) + Eventually(started, 5*time.Second).Should(BeClosed()) + stop() + Expect(sub.Unsubscribe()).To(Succeed()) + + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed()) + Expect(row.State).To(Equal(ClaimPending), "released, not left held by a replica that is leaving") + }) + + It("refuses a kind that is not a work kind", func() { + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus}) + _, err := c.Consume(ctx, messaging.WorkKind("nonsense"), 0, func(context.Context, []byte, messaging.Publisher) error { return nil }) + Expect(err).To(HaveOccurred()) + }) + + It("is built only with what it needs", func() { + good := ClaimConsumerConfig{DB: db, Owner: "replica-a", Events: bus} + for name, mutate := range map[string]func(*ClaimConsumerConfig){ + "no database": func(c *ClaimConsumerConfig) { c.DB = nil }, + "no owner": func(c *ClaimConsumerConfig) { c.Owner = "" }, + "no events": func(c *ClaimConsumerConfig) { c.Events = nil }, + } { + cfg := good + mutate(&cfg) + _, err := NewClaimConsumer(cfg) + Expect(err).To(HaveOccurred(), name) + } + _, err := NewClaimConsumer(good) + Expect(err).ToNot(HaveOccurred()) + }) + }) + + Describe("the order of the kinds in the table", func() { + It("serves a kind that is the only one with work, whatever else is waiting", func() { + queue := NewClaimQueue(db, bus) + for i := range 3 { + Expect(queue.Enqueue(ctx, messaging.WorkTask, i)).To(Succeed()) + } + var got atomic.Int32 + c := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 50 * time.Millisecond}) + consume(c, messaging.WorkAgentRun, 0, func(context.Context, []byte, messaging.Publisher) error { got.Add(1); return nil }) + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, 1)).To(Succeed()) + Eventually(got.Load, 5*time.Second).Should(BeEquivalentTo(1)) + Expect(rows()).To(BeEquivalentTo(3)) + }) + }) +}) + diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 65a53d99f1e5..c646fa9acd62 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -20,6 +20,12 @@ const ( QueueWorkers = "workers" ) +// SubjectClaimWake is the hint that a unit of work was put in the claim queue. +// A consumer that hears it looks for work at once and does not wait for its next +// poll. The hint is a broadcast and broadcasts are at-most-once, so it can be +// lost, and the poll finds the work then. +const SubjectClaimWake = "jobs.claim.wake" + // Status Updates (Pub/Sub — all subscribers get every message, for SSE bridging) // These use parameterized subjects: e.g. SubjectAgentEvents("myagent", "user1") const ( From d1e50b153775fb174a0866dd908605405133d4e2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:21:13 +0000 Subject: [PATCH 081/162] refactor(nodes): let the claim loop pick an agent worker and tell a busy one The loop that drives claims lives in package jobs, which does not import nodes. It needs two things from nodes: to pick a worker while leaving out the ones already tried, and to tell a worker with no free slot from a worker that failed. The pick takes the exclusion set that the agent control already used, and the sentinel for a busy worker moves to workerctl, where both packages can reach it. nodes.ErrWorkerBusy stays as the same value. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/agent_control.go | 2 +- core/services/nodes/agent_control_test.go | 18 +++++++++--------- core/services/nodes/agent_selector.go | 19 ++++++++----------- core/services/nodes/control_client.go | 8 +++----- core/services/workerctl/request.go | 7 +++++++ 5 files changed, 28 insertions(+), 26 deletions(-) diff --git a/core/services/nodes/agent_control.go b/core/services/nodes/agent_control.go index 3af7ac3f6795..5e21cbcd0a4a 100644 --- a/core/services/nodes/agent_control.go +++ b/core/services/nodes/agent_control.go @@ -80,7 +80,7 @@ func agentVerb[Req, Rep any](ctx context.Context, a *AgentControlClient, verb st tried := make(map[string]bool, maxAgentPicks) var last error for range maxAgentPicks { - nodeID, _, err := a.sel.pickConnectedExcluding(callCtx, tried) + nodeID, _, err := a.sel.PickConnectedExcluding(callCtx, tried) if err != nil { if last != nil && errors.Is(err, ErrNoAgentWorker) { // Every worker that holds a tunnel has failed to take the request. diff --git a/core/services/nodes/agent_control_test.go b/core/services/nodes/agent_control_test.go index 1a2820835076..eff7a09e55e8 100644 --- a/core/services/nodes/agent_control_test.go +++ b/core/services/nodes/agent_control_test.go @@ -146,21 +146,21 @@ var _ = Describe("The agent control of the tunnel carrier", func() { Describe("choosing an agent worker", func() { It("reports no agent worker, as a missing route, when none is registered", func() { - _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).To(MatchError(ErrNoAgentWorker)) Expect(errors.Is(err, ErrNoRoute)).To(BeTrue()) }) It("reports no agent worker when none of the registered ones holds a tunnel", func() { register("idle", NodeTypeAgent, true) - _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).To(MatchError(ErrNoAgentWorker)) }) It("keeps a database error apart from a missing route", func() { register("a1", NodeTypeAgent, true) conns.err = errors.New("database is away") - _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).To(HaveOccurred()) Expect(errors.Is(err, ErrNoAgentWorker)).To(BeFalse()) Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) @@ -171,7 +171,7 @@ var _ = Describe("The agent control of the tunnel carrier", func() { own := register("own", NodeTypeAgent, true) conns.held[remote], conns.held[own] = false, true for range 20 { - id, nodeType, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + id, nodeType, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).ToNot(HaveOccurred()) Expect(id).To(Equal(own)) Expect(nodeType).To(Equal(NodeTypeAgent)) @@ -181,7 +181,7 @@ var _ = Describe("The agent control of the tunnel carrier", func() { It("falls back to a worker that another replica holds", func() { remote := register("remote", NodeTypeAgent, true) conns.held[remote] = false - id, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + id, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).ToNot(HaveOccurred()) Expect(id).To(Equal(remote)) }) @@ -194,7 +194,7 @@ var _ = Describe("The agent control of the tunnel carrier", func() { for _, id := range []string{backend, pending, draining} { conns.held[id] = true } - _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + _, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).To(MatchError(ErrNoAgentWorker)) Expect(conns.asked).ToNot(BeEmpty()) for _, asked := range conns.asked { @@ -206,7 +206,7 @@ var _ = Describe("The agent control of the tunnel carrier", func() { id := register("flaky", NodeTypeAgent, true) Expect(reg.MarkUnhealthy(ctx, id)).To(Succeed()) conns.held[id] = true - picked, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnected(ctx) + picked, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, nil) Expect(err).ToNot(HaveOccurred()) Expect(picked).To(Equal(id)) }) @@ -215,14 +215,14 @@ var _ = Describe("The agent control of the tunnel carrier", func() { a := register("a", NodeTypeAgent, true) b := register("b", NodeTypeAgent, true) conns.held[a], conns.held[b] = true, true - picked, _, err := NewAgentSelector(reg, conns, "replica-a").pickConnectedExcluding(ctx, map[string]bool{a: true}) + picked, _, err := NewAgentSelector(reg, conns, "replica-a").PickConnectedExcluding(ctx, map[string]bool{a: true}) Expect(err).ToNot(HaveOccurred()) Expect(picked).To(Equal(b)) }) It("answers as a missing route when it was built with nothing", func() { var sel *AgentSelector - _, _, err := sel.PickConnected(ctx) + _, _, err := sel.PickConnectedExcluding(ctx, nil) Expect(err).To(MatchError(ErrNoAgentWorker)) }) }) diff --git a/core/services/nodes/agent_selector.go b/core/services/nodes/agent_selector.go index dae1c8953e4f..3d26c552ba79 100644 --- a/core/services/nodes/agent_selector.go +++ b/core/services/nodes/agent_selector.go @@ -46,17 +46,14 @@ func NewAgentSelector(registry *NodeRegistry, conns AgentConnectionReader, self return &AgentSelector{registry: registry, conns: conns, self: self} } -// PickConnected returns the id and the node type of an agent node whose tunnel a -// live replica holds. The type comes back because the caller that publishes the -// events of a run needs it for every line, and the selector has read the node -// rows already. -func (s *AgentSelector) PickConnected(ctx context.Context) (nodeID, nodeType string, err error) { - return s.pickConnectedExcluding(ctx, nil) -} - -// pickConnectedExcluding is PickConnected without the nodes in tried. A retry -// that could pick the worker that has just failed would be the same call again. -func (s *AgentSelector) pickConnectedExcluding(ctx context.Context, tried map[string]bool) (string, string, error) { +// PickConnectedExcluding returns the id and the node type of an agent node whose +// tunnel a live replica holds, leaving out the nodes in tried. The type comes back +// because the caller that publishes the events of a run needs it for every line, +// and the selector has read the node rows already. +// +// The exclusion is for a retry: a retry that could pick the worker that has just +// failed would be the same call again. +func (s *AgentSelector) PickConnectedExcluding(ctx context.Context, tried map[string]bool) (string, string, error) { if s == nil || s.registry == nil || s.conns == nil { return "", "", fmt.Errorf("this deployment has no agent selector: %w", ErrNoAgentWorker) } diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go index 66b537665f87..7f90b2bf620f 100644 --- a/core/services/nodes/control_client.go +++ b/core/services/nodes/control_client.go @@ -238,11 +238,9 @@ func (c *ControlClient) do(ctx context.Context, nodeID, verb string, req any) (* } } -// ErrWorkerBusy means that the worker answered that it has no free slot for a -// run. The worker is present and the request did not start, so it says nothing -// about the work, and a caller may offer the run to another worker. It is not -// ErrNoRoute: the route exists. -var ErrWorkerBusy = errors.New("the worker has no free slot") +// ErrWorkerBusy is workerctl.ErrWorkerBusy: the worker answered that it has no +// free slot for a run. +var ErrWorkerBusy = workerctl.ErrWorkerBusy // maxControlErrorBodyBytes bounds how much of an error body reaches a log line. const maxControlErrorBodyBytes = 512 diff --git a/core/services/workerctl/request.go b/core/services/workerctl/request.go index dbd9e6fbba52..52466c4f3ddf 100644 --- a/core/services/workerctl/request.go +++ b/core/services/workerctl/request.go @@ -1,6 +1,7 @@ package workerctl import ( + "errors" "io" "net/http" "strings" @@ -59,6 +60,12 @@ func WriteUnknownPath(w http.ResponseWriter, r *http.Request) { // full from a worker that failed. const BusyMessage = "the worker has no free slot for this run" +// ErrWorkerBusy means that a worker answered that it has no free slot for a run. +// The worker is present and the request did not start, so it says nothing about +// the work, and a caller may offer the run to another worker. It is not a missing +// route: the route exists. +var ErrWorkerBusy = errors.New("the worker has no free slot") + // WriteBusy answers a run that this worker cannot take now. The worker is // present and routable, and the request was not started, so the frontend may // offer it to another worker. From 81d60baccccf10cf661cc059b26752472b2ddab2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:21:13 +0000 Subject: [PATCH 082/162] feat(jobs): drive claims on agent workers and end a run when its job is cancelled A frontend replica claims the work and an agent worker runs it, because an agent worker has no database and cannot claim. The AgentDriver is the handler of the claims: it picks a worker, sends the payload to the verb of its kind and reads the stream that comes back. The lines that name a subject are published through the rebroadcaster, and the last line is the reply. The terminal state of a job is written before the claim is completed, by the order of the calls. If the write is refused the claim is kept with its replica, because the work ran and releasing it would run it again. A call that obtained no answer releases the claim with a wait. A worker with no free slot is not a failure: the run is offered to the next worker, up to three. Dispatcher.Cancel publishes jobs..cancel and nothing subscribed to it on master. On this carrier the replica that holds the job ends its stream, the run on the worker ends with its request, and the job is closed as cancelled. The NATS carrier stays as it was, and a spec pins that. A plain task has no worker on either carrier. It is dropped here and the job is left to the reaper, as it was on NATS. DispatchLoop is a ClaimConsumer for the three kinds with the driver as its handler, four runs at a time for each kind. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/jobs/claim_driver.go | 292 ++++++++++++++++++ core/services/jobs/claim_driver_test.go | 359 +++++++++++++++++++++++ core/services/jobs/claim_queue_test.go | 7 +- core/services/jobs/dispatch_loop.go | 118 ++++++++ core/services/jobs/dispatch_loop_test.go | 177 +++++++++++ core/services/jobs/dispatcher_test.go | 17 ++ core/services/messaging/subjects.go | 1 + core/services/testutil/fakebus.go | 12 + 8 files changed, 981 insertions(+), 2 deletions(-) create mode 100644 core/services/jobs/claim_driver.go create mode 100644 core/services/jobs/claim_driver_test.go create mode 100644 core/services/jobs/dispatch_loop.go create mode 100644 core/services/jobs/dispatch_loop_test.go diff --git a/core/services/jobs/claim_driver.go b/core/services/jobs/claim_driver.go new file mode 100644 index 000000000000..d502641b53d8 --- /dev/null +++ b/core/services/jobs/claim_driver.go @@ -0,0 +1,292 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "sync/atomic" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// maxDrivePicks bounds how many agent workers one run is offered to. A worker +// that is full is no fault of the fleet, but a fleet whose first three workers +// have no slot is busy, and the claim goes back to the pool with a wait. +const maxDrivePicks = 3 + +// The ports below exist because neither jobs to nodes nor nodes to jobs is an +// import edge, and neither should become one for a driver. Each is named for what +// it does here and not for the type that satisfies it. + +// AgentPicker chooses an agent worker for a run, leaving out those already tried. +// Satisfied by *nodes.AgentSelector. +type AgentPicker interface { + PickConnectedExcluding(ctx context.Context, tried map[string]bool) (nodeID, nodeType string, err error) +} + +// StreamCaller issues one streaming control request. Satisfied by +// *nodes.ControlClient. +type StreamCaller interface { + CallStreaming(ctx context.Context, nodeID, verb string, req, reply any, + onProgress func(subject string, raw json.RawMessage)) error +} + +// ProgressBroadcaster publishes the line that a worker asked to have published, +// when the node type of the worker may ask for it. Satisfied by +// *nodes.Rebroadcaster. It returns a bool and no error on purpose: a refused or +// failed publish is not the verdict of the worker about the work. +type ProgressBroadcaster interface { + Handle(nodeType, subject string, raw json.RawMessage) bool +} + +// TerminalWriter records the terminal state of a job. Satisfied by *JobStore. +type TerminalWriter interface { + UpdateJobStatus(jobID, status, result, errMsg string) error +} + +var _ TerminalWriter = (*JobStore)(nil) + +// AgentDriverConfig is what an AgentDriver needs. +type AgentDriverConfig struct { + Picker AgentPicker + Control StreamCaller + Broadcast ProgressBroadcaster + // Store records the terminal state of the jobs that name one. Without it a + // job is not closed by this driver, and its result reaches the store only + // through the broadcast of the worker. + Store TerminalWriter + // Hints is where the cancel of a job is heard. Without it a cancel does not + // reach a run. + Hints messaging.Broadcaster +} + +// AgentDriver is the handler of the claims on the tunnel carrier. A frontend +// replica claims the work, and the driver hands it to an agent worker and reads +// the stream that comes back. The replica claims, and the worker does not, +// because an agent worker has no database. +// +// The events of the run, its progress and its result arrive as lines of the +// response that this replica is already reading. That is what lets the terminal +// state be written before the claim is completed, by the order of the calls and +// not by a retry. +type AgentDriver struct { + cfg AgentDriverConfig + + // cancels holds a function that ends the stream of each job that this replica + // drives. A cancel is heard by every replica and applied by the one that holds + // the job. + cancels messaging.CancelRegistry + // cancelled remembers the jobs whose stream a cancel ended, so that the end + // of the stream is read as the answer to the cancel and not as a link that + // broke. + cancelled cancelledJobs + sub atomic.Pointer[messaging.Subscription] +} + +// NewAgentDriver returns a driver. It refuses a configuration that would claim +// work and never hand it on. +func NewAgentDriver(cfg AgentDriverConfig) (*AgentDriver, error) { + if cfg.Picker == nil { + return nil, errors.New("the agent driver was built with no way to pick an agent worker") + } + if cfg.Control == nil { + return nil, errors.New("the agent driver was built with no control transport to reach an agent worker over") + } + return &AgentDriver{cfg: cfg}, nil +} + +// Start listens for the cancel of jobs. +// +// A cancel has no consumer on the NATS carrier: Dispatcher.Cancel publishes it +// and nothing subscribes, as before. On this carrier the replica that holds the +// job ends its stream, and the run on the worker ends with its request. +func (d *AgentDriver) Start(_ context.Context) error { + if d.cfg.Hints == nil { + return nil + } + sub, err := messaging.SubscribeJSON(d.cfg.Hints, messaging.SubjectJobCancelWildcard, func(evt CancelEvent) { + d.cancelJob(evt.JobID) + }) + if err != nil { + return fmt.Errorf("subscribing to the cancel of jobs: %w", err) + } + d.sub.Store(&sub) + return nil +} + +// Stop ends the subscription to the cancel of jobs. +func (d *AgentDriver) Stop() { + if sub := d.sub.Swap(nil); sub != nil { + _ = (*sub).Unsubscribe() + } +} + +func (d *AgentDriver) cancelJob(jobID string) { + if jobID == "" { + return + } + // Marked first: the stream ends as soon as the function runs, and the driver + // that sees it end must already know why. + d.cancelled.add(jobID) + if !d.cancels.Cancel(jobID) { + d.cancelled.remove(jobID) + } +} + +// Handler returns the handler of a kind, for a ClaimConsumer. +func (d *AgentDriver) Handler(kind messaging.WorkKind) (messaging.WorkHandler, error) { + switch kind { + case messaging.WorkAgentRun: + return d.handler(kind, workerctl.VerbAgentExecute), nil + case messaging.WorkMCPCI: + return d.handler(kind, workerctl.VerbMCPCIRun), nil + case messaging.WorkTask: + return d.task, nil + } + return nil, fmt.Errorf("the agent driver has no handler for the work kind %q", kind) +} + +// task answers a plain task. No worker serves one, and none did on the NATS +// carrier: the job was published to a subject nobody consumed and stayed running +// until the reaper failed it. The claim row makes that visible and does not +// change it. The row is completed, so that it does not wait for ever, and the job +// is left to the same reaper, so that both carriers behave alike. +func (d *AgentDriver) task(_ context.Context, payload []byte, _ messaging.Publisher) error { + xlog.Warn("Dropping a plain task job: no worker in this deployment serves it, and the reaper will fail the job", + "job", jobIDOf(payload)) + return nil +} + +func (d *AgentDriver) handler(kind messaging.WorkKind, verb string) messaging.WorkHandler { + return func(ctx context.Context, payload []byte, _ messaging.Publisher) error { + return d.drive(ctx, kind, verb, payload) + } +} + +// drive hands one claim to a worker and returns what settles it: nil when the +// worker answered (or the job was cancelled), ErrKeepClaim when the worker +// answered and the answer could not be written, and any other error when nothing +// was learned about the work. +func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb string, payload []byte) error { + jobID := "" + if kind == messaging.WorkMCPCI { + jobID = jobIDOf(payload) + } + if jobID != "" { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(ctx) + defer cancel() + d.cancels.Register(jobID, cancel) + defer d.cancels.Deregister(jobID) + defer d.cancelled.remove(jobID) + } + + tried := make(map[string]bool, maxDrivePicks) + var last error + for range maxDrivePicks { + nodeID, nodeType, err := d.cfg.Picker.PickConnectedExcluding(ctx, tried) + if err != nil { + if last != nil { + // Every worker that was offered the run had no slot. That is the + // evidence, and the end of the list is not. + return last + } + return fmt.Errorf("picking an agent worker for a %s claim: %w", kind, err) + } + tried[nodeID] = true + + var reply workerctl.RunReply + err = d.cfg.Control.CallStreaming(ctx, nodeID, verb, json.RawMessage(payload), &reply, + func(subject string, raw json.RawMessage) { d.rebroadcast(nodeType, subject, raw) }) + switch { + case err == nil: + return d.persistTerminal(&reply) + case errors.Is(err, workerctl.ErrWorkerBusy): + last = fmt.Errorf("offering a %s claim to agent worker %q: %w", kind, nodeID, err) + continue + case jobID != "" && d.cancelled.has(jobID) && ctx.Err() != nil: + // The stream was ended by the cancel of the job. The worker's run dies + // with its request, and the job is closed here, because nothing else + // will. + return d.persistTerminal(&workerctl.RunReply{JobID: jobID, Status: "cancelled", Error: "cancelled"}) + default: + return fmt.Errorf("driving a %s claim on agent worker %q: %w", kind, nodeID, err) + } + } + return last +} + +// rebroadcast publishes the broadcast that a progress line asked for. A line with +// no subject is meant for the caller alone, as every private tick is, and is not +// passed on: the allow list denies the empty subject, and passing it would log a +// refusal for each ordinary tick. +func (d *AgentDriver) rebroadcast(nodeType, subject string, raw json.RawMessage) { + if subject == "" || d.cfg.Broadcast == nil { + return + } + d.cfg.Broadcast.Handle(nodeType, subject, raw) +} + +// persistTerminal writes the terminal state of the job that a reply names. A run +// that names none is an agent chat, whose output reaches the user as events. +// +// A reply that names a job and no status closes the job as failed. The outcome +// to avoid is a job left running after its work has ended. +func (d *AgentDriver) persistTerminal(reply *workerctl.RunReply) error { + if reply == nil || reply.JobID == "" || d.cfg.Store == nil { + return nil + } + status, errMsg := reply.Status, reply.Error + if status == "" { + status, errMsg = "failed", "the worker answered without naming a terminal status" + } + if err := d.cfg.Store.UpdateJobStatus(reply.JobID, status, reply.Result, errMsg); err != nil { + // The work ran. Releasing the claim would run it again, so the claim stays + // with this replica, and it becomes claimable only if this replica dies. + return fmt.Errorf("writing the terminal state of job %q: %w: %w", reply.JobID, err, ErrKeepClaim) + } + return nil +} + +// jobIDOf reads the job id out of the payload of an MCP CI claim. It returns "" +// for anything that is not a JobEvent. +func jobIDOf(payload []byte) string { + var evt JobEvent + if err := json.Unmarshal(payload, &evt); err != nil { + return "" + } + return evt.JobID +} + +// cancelledJobs is a set of job ids, safe for concurrent use. +type cancelledJobs struct { + mu sync.Mutex + ids map[string]struct{} +} + +func (c *cancelledJobs) add(id string) { + c.mu.Lock() + defer c.mu.Unlock() + if c.ids == nil { + c.ids = map[string]struct{}{} + } + c.ids[id] = struct{}{} +} + +func (c *cancelledJobs) remove(id string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.ids, id) +} + +func (c *cancelledJobs) has(id string) bool { + c.mu.Lock() + defer c.mu.Unlock() + _, ok := c.ids[id] + return ok +} diff --git a/core/services/jobs/claim_driver_test.go b/core/services/jobs/claim_driver_test.go new file mode 100644 index 000000000000..b9364ddbd67a --- /dev/null +++ b/core/services/jobs/claim_driver_test.go @@ -0,0 +1,359 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// fakePicker hands out a scripted sequence of workers. +type fakePicker struct { + mu sync.Mutex + nodes []string + types map[string]string + err error + asked [][]string +} + +func (p *fakePicker) PickConnectedExcluding(_ context.Context, tried map[string]bool) (string, string, error) { + p.mu.Lock() + defer p.mu.Unlock() + var skipped []string + for id := range tried { + skipped = append(skipped, id) + } + p.asked = append(p.asked, skipped) + if p.err != nil { + return "", "", p.err + } + for _, id := range p.nodes { + if !tried[id] { + return id, p.types[id], nil + } + } + return "", "", errNoWorkerLeft +} + +var errNoWorkerLeft = errors.New("no worker left") + +type call struct { + node, verb string + payload string +} + +// fakeControl plays the control client. do decides what each call does. +type fakeControl struct { + mu sync.Mutex + calls []call + do func(ctx context.Context, node, verb string, onProgress func(string, json.RawMessage)) (workerctl.RunReply, error) +} + +func (f *fakeControl) CallStreaming(ctx context.Context, node, verb string, req, reply any, onProgress func(string, json.RawMessage)) error { + raw, _ := req.(json.RawMessage) + f.mu.Lock() + f.calls = append(f.calls, call{node, verb, string(raw)}) + do := f.do + f.mu.Unlock() + r, err := do(ctx, node, verb, onProgress) + if err != nil { + return err + } + *(reply.(*workerctl.RunReply)) = r + return nil +} + +func (f *fakeControl) seen() []call { + f.mu.Lock() + defer f.mu.Unlock() + return append([]call(nil), f.calls...) +} + +type progressLine struct{ nodeType, subject, raw string } + +type fakeBroadcast struct { + mu sync.Mutex + lines []progressLine +} + +func (b *fakeBroadcast) Handle(nodeType, subject string, raw json.RawMessage) bool { + b.mu.Lock() + defer b.mu.Unlock() + b.lines = append(b.lines, progressLine{nodeType, subject, string(raw)}) + return true +} + +type terminal struct{ jobID, status, result, errMsg string } + +type fakeStore struct { + mu sync.Mutex + got []terminal + fail error +} + +func (s *fakeStore) UpdateJobStatus(jobID, status, result, errMsg string) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.fail != nil { + return s.fail + } + s.got = append(s.got, terminal{jobID, status, result, errMsg}) + return nil +} + +func (s *fakeStore) terminals() []terminal { + s.mu.Lock() + defer s.mu.Unlock() + return append([]terminal(nil), s.got...) +} + +var _ = Describe("The agent driver", func() { + var ( + ctx context.Context + picker *fakePicker + control *fakeControl + broadcast *fakeBroadcast + store *fakeStore + bus *testutil.FakeBus + driver *AgentDriver + ) + + ciPayload := func(jobID string) []byte { + raw, err := json.Marshal(JobEvent{JobID: jobID, TaskID: "t1"}) + Expect(err).ToNot(HaveOccurred()) + return raw + } + + BeforeEach(func() { + c, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + ctx = c + picker = &fakePicker{nodes: []string{"w1", "w2", "w3", "w4"}, types: map[string]string{"w1": "agent", "w2": "agent", "w3": "agent", "w4": "agent"}} + control = &fakeControl{do: func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, nil + }} + broadcast = &fakeBroadcast{} + store = &fakeStore{} + bus = testutil.NewFakeBus() + var err error + driver, err = NewAgentDriver(AgentDriverConfig{Picker: picker, Control: control, Broadcast: broadcast, Store: store, Hints: bus}) + Expect(err).ToNot(HaveOccurred()) + Expect(driver.Start(ctx)).To(Succeed()) + DeferCleanup(driver.Stop) + }) + + handler := func(kind messaging.WorkKind) messaging.WorkHandler { + GinkgoHelper() + h, err := driver.Handler(kind) + Expect(err).ToNot(HaveOccurred()) + return h + } + + Describe("a run", func() { + It("sends the payload to the verb of its kind on the worker the picker chose", func() { + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{"agent_name":"a"}`), nil)).To(Succeed()) + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed()) + calls := control.seen() + Expect(calls).To(HaveLen(2)) + Expect(calls[0]).To(Equal(call{"w1", workerctl.VerbAgentExecute, `{"agent_name":"a"}`})) + Expect(calls[1].verb).To(Equal(workerctl.VerbMCPCIRun)) + Expect(calls[1].payload).To(MatchJSON(string(ciPayload("j1")))) + }) + + It("publishes the lines that name a subject, with the node type of the worker, and drops the others", func() { + control.do = func(_ context.Context, _, _ string, onProgress func(string, json.RawMessage)) (workerctl.RunReply, error) { + onProgress("agent.a.events.u", json.RawMessage(`{"n":1}`)) + onProgress("", json.RawMessage(`{"private":"tick"}`)) + onProgress("jobs.j1.progress", json.RawMessage(`{"n":2}`)) + return workerctl.RunReply{}, nil + } + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).To(Succeed()) + Expect(broadcast.lines).To(Equal([]progressLine{ + {"agent", "agent.a.events.u", `{"n":1}`}, + {"agent", "jobs.j1.progress", `{"n":2}`}, + })) + }) + + It("writes the terminal state of a job before it lets the claim be completed", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{JobID: "j1", Status: "completed", Result: "42"}, nil + } + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed()) + Expect(store.terminals()).To(Equal([]terminal{{"j1", "completed", "42", ""}})) + }) + + It("closes a job whose reply names no status as failed, because the one outcome to avoid is a job left running", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{JobID: "j1"}, nil + } + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed()) + got := store.terminals() + Expect(got).To(HaveLen(1)) + Expect(got[0].status).To(Equal("failed")) + Expect(got[0].errMsg).ToNot(BeEmpty()) + }) + + It("writes no job for a run that names none", func() { + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).To(Succeed()) + Expect(store.terminals()).To(BeEmpty()) + }) + + It("keeps the claim, and does not release it, when the terminal state cannot be written", func() { + store.fail = errors.New("database is away") + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{JobID: "j1", Status: "completed"}, nil + } + err := handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil) + Expect(errors.Is(err, ErrKeepClaim)).To(BeTrue(), "the work ran, so releasing would run it twice; got %v", err) + }) + + It("returns the failure of a call that obtained no answer, so that the claim is released", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, fmt.Errorf("the stream ended before its reply line: %w", errors.New("unexpected EOF")) + } + err := handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrKeepClaim)).To(BeFalse()) + Expect(store.terminals()).To(BeEmpty(), "a broken link is not a verdict about the job") + }) + + It("returns the error of the picker when no worker is connected", func() { + picker.err = errors.New("no agent worker holds a tunnel") + err := handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil) + Expect(err).To(MatchError(ContainSubstring("no agent worker"))) + Expect(control.seen()).To(BeEmpty()) + }) + + It("offers the run to the next worker when the first has no free slot, up to three", func() { + control.do = func(_ context.Context, node, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + if node == "w3" { + return workerctl.RunReply{}, nil + } + return workerctl.RunReply{}, fmt.Errorf("control request: %w", workerctl.ErrWorkerBusy) + } + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).To(Succeed()) + Expect(control.seen()).To(HaveLen(3)) + + control.mu.Lock() + control.calls = nil + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, workerctl.ErrWorkerBusy + } + control.mu.Unlock() + err := handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil) + Expect(errors.Is(err, workerctl.ErrWorkerBusy)).To(BeTrue()) + Expect(control.seen()).To(HaveLen(3), "a fourth worker is not tried") + }) + + It("does not offer a run to a second worker after a failure that may have started it", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, errors.New("the stream ended before its reply line") + } + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).ToNot(Succeed()) + Expect(control.seen()).To(HaveLen(1)) + }) + + It("leaves a plain task alone, as the NATS carrier does: nothing serves it and the job stays as it is", func() { + Expect(handler(messaging.WorkTask)(ctx, ciPayload("j1"), nil)).To(Succeed()) + Expect(control.seen()).To(BeEmpty()) + Expect(store.terminals()).To(BeEmpty()) + }) + + It("refuses a kind it does not know", func() { + _, err := driver.Handler(messaging.WorkKind("nonsense")) + Expect(err).To(HaveOccurred()) + }) + }) + + Describe("the cancel of a job", func() { + running := func() (started chan struct{}, finished chan error) { + started = make(chan struct{}) + finished = make(chan error, 1) + control.do = func(ctx context.Context, _, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + close(started) + <-ctx.Done() + return workerctl.RunReply{}, ctx.Err() + } + go func() { finished <- handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil) }() + return started, finished + } + + It("ends the stream of the job, so the run on the worker dies with its request, and records the job as cancelled", func() { + started, finished := running() + Eventually(started, 5*time.Second).Should(BeClosed()) + + Expect(bus.Publish(messaging.SubjectJobCancel("j1"), CancelEvent{JobID: "j1"})).To(Succeed()) + + var err error + Eventually(finished, 5*time.Second).Should(Receive(&err)) + Expect(err).ToNot(HaveOccurred(), "the cancel is an answer, so the claim is completed and not released") + Expect(store.terminals()).To(HaveLen(1)) + Expect(store.terminals()[0].jobID).To(Equal("j1")) + Expect(store.terminals()[0].status).To(Equal("cancelled")) + }) + + It("ignores the cancel of a job that runs elsewhere or nowhere", func() { + started, finished := running() + Eventually(started, 5*time.Second).Should(BeClosed()) + Expect(bus.Publish(messaging.SubjectJobCancel("other"), CancelEvent{JobID: "other"})).To(Succeed()) + Consistently(finished, 300*time.Millisecond).ShouldNot(Receive()) + Expect(store.terminals()).To(BeEmpty()) + // Let the run go: the context of the spec ends it. + }) + + It("treats the stop of the replica as a failure to answer and not as a cancel", func() { + started := make(chan struct{}) + control.do = func(ctx context.Context, _, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + close(started) + <-ctx.Done() + return workerctl.RunReply{}, ctx.Err() + } + parent, stop := context.WithCancel(ctx) + finished := make(chan error, 1) + go func() { finished <- handler(messaging.WorkMCPCI)(parent, ciPayload("j1"), nil) }() + Eventually(started, 5*time.Second).Should(BeClosed()) + stop() + var err error + Eventually(finished, 5*time.Second).Should(Receive(&err)) + Expect(err).To(HaveOccurred()) + Expect(store.terminals()).To(BeEmpty()) + }) + + It("stops hearing cancels after Stop", func() { + started, finished := running() + Eventually(started, 5*time.Second).Should(BeClosed()) + driver.Stop() + Expect(bus.Publish(messaging.SubjectJobCancel("j1"), CancelEvent{JobID: "j1"})).To(Succeed()) + Consistently(finished, 300*time.Millisecond).ShouldNot(Receive()) + }) + + It("does not leave a cancel function behind for a run that ended", func() { + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed()) + Expect(driver.cancels.Cancel("j1")).To(BeFalse()) + }) + }) + + It("is built only with what it needs", func() { + good := AgentDriverConfig{Picker: picker, Control: control, Broadcast: broadcast, Store: store, Hints: bus} + for name, mutate := range map[string]func(*AgentDriverConfig){ + "no picker": func(c *AgentDriverConfig) { c.Picker = nil }, + "no control": func(c *AgentDriverConfig) { c.Control = nil }, + } { + cfg := good + mutate(&cfg) + _, err := NewAgentDriver(cfg) + Expect(err).To(HaveOccurred(), name) + } + _, err := NewAgentDriver(AgentDriverConfig{Picker: picker, Control: control}) + Expect(err).ToNot(HaveOccurred(), "a driver with no store and no bus still drives; it records and cancels nothing") + }) +}) diff --git a/core/services/jobs/claim_queue_test.go b/core/services/jobs/claim_queue_test.go index 5b1a22376eab..840f5c27fa3e 100644 --- a/core/services/jobs/claim_queue_test.go +++ b/core/services/jobs/claim_queue_test.go @@ -261,7 +261,11 @@ var _ = Describe("The claim queue", func() { return nil }) } - Eventually(func() int { mu.Lock(); defer mu.Unlock(); return len(got[messaging.WorkMCPCI]) + len(got[messaging.WorkAgentRun]) }, 15*time.Second).Should(Equal(12)) + Eventually(func() int { + mu.Lock() + defer mu.Unlock() + return len(got[messaging.WorkMCPCI]) + len(got[messaging.WorkAgentRun]) + }, 15*time.Second).Should(Equal(12)) mu.Lock() defer mu.Unlock() Expect(got[messaging.WorkMCPCI]).To(Equal([]int{0, 1, 2, 3, 4, 5})) @@ -582,4 +586,3 @@ var _ = Describe("The claim queue", func() { }) }) }) - diff --git a/core/services/jobs/dispatch_loop.go b/core/services/jobs/dispatch_loop.go new file mode 100644 index 000000000000..dc9790051568 --- /dev/null +++ b/core/services/jobs/dispatch_loop.go @@ -0,0 +1,118 @@ +package jobs + +import ( + "context" + "errors" + "fmt" + "time" + + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// DefaultMaxConcurrent is how many claims of one kind a replica drives at once. +// +// A claim is driven while the response of the worker is read: the progress, the +// events and the terminal line all arrive on it, and an MCP CI job holds it for +// up to ten minutes. With one slot, the second queued job of a deployment would +// wait for the first. The limit is for each kind, so a long job of one kind does +// not keep another kind waiting. +const DefaultMaxConcurrent = 4 + +// DispatchConfig is everything a DispatchLoop needs. +type DispatchConfig struct { + DB *gorm.DB + // Owner is the id of this replica in the instances table. + Owner string + Picker AgentPicker + Control StreamCaller + Broadcast ProgressBroadcaster + Store TerminalWriter + // Hints carries the wake hints of the queue, the cancel of jobs and, for the + // handlers, the place for events. + Hints messaging.Broadcaster + // Interval is the poll, DefaultClaimInterval when zero. + Interval time.Duration + // Liveness is the window of replica liveness, cluster.InstanceLiveness when + // zero. + Liveness time.Duration + // MaxConcurrent is how many claims of a kind this replica drives at once, + // DefaultMaxConcurrent when zero. + MaxConcurrent int +} + +// DispatchLoop is what a frontend replica runs on the tunnel carrier: it claims +// the queued work of every kind and drives it on agent workers. It is a +// ClaimConsumer whose handlers are the AgentDriver. +type DispatchLoop struct { + driver *AgentDriver + consumer *ClaimConsumer + max int + subs []messaging.Subscription +} + +// NewDispatchLoop returns the loop of one replica. Every refusal is a piece of +// wiring whose absence has no symptom: a loop with no picker claims rows and +// releases them for ever, and one with no owner writes claims that the reap +// cannot attribute. Both look like work that is accepted and never done, which +// cannot be told at the API from a busy deployment. +func NewDispatchLoop(cfg DispatchConfig) (*DispatchLoop, error) { + if cfg.Hints == nil { + return nil, errors.New("the dispatch loop was built with no bus for wake hints and cancels") + } + driver, err := NewAgentDriver(AgentDriverConfig{ + Picker: cfg.Picker, Control: cfg.Control, Broadcast: cfg.Broadcast, Store: cfg.Store, Hints: cfg.Hints, + }) + if err != nil { + return nil, err + } + consumer, err := NewClaimConsumer(ClaimConsumerConfig{ + DB: cfg.DB, Owner: cfg.Owner, Events: cfg.Hints, Hints: cfg.Hints, + Interval: cfg.Interval, Liveness: cfg.Liveness, + }) + if err != nil { + return nil, err + } + max := cfg.MaxConcurrent + if max <= 0 { + max = DefaultMaxConcurrent + } + return &DispatchLoop{driver: driver, consumer: consumer, max: max}, nil +} + +// Start begins to claim and drive until ctx ends or Stop is called. +func (l *DispatchLoop) Start(ctx context.Context) error { + if err := l.driver.Start(ctx); err != nil { + return err + } + for _, kind := range []messaging.WorkKind{messaging.WorkMCPCI, messaging.WorkAgentRun, messaging.WorkTask} { + h, err := l.driver.Handler(kind) + if err != nil { + l.Stop() + return err + } + sub, err := l.consumer.Consume(ctx, kind, l.max, h) + if err != nil { + l.Stop() + return fmt.Errorf("consuming the %s claims: %w", kind, err) + } + l.subs = append(l.subs, sub) + } + return nil +} + +// Stop stops claiming and waits until every run in flight has been settled. A +// claim that a process leaves held can be released by no other replica until the +// liveness window has passed, so Stop does not return before the claims are +// settled. +func (l *DispatchLoop) Stop() { + if l == nil { + return + } + for _, sub := range l.subs { + _ = sub.Unsubscribe() + } + l.subs = nil + l.driver.Stop() +} diff --git a/core/services/jobs/dispatch_loop_test.go b/core/services/jobs/dispatch_loop_test.go new file mode 100644 index 000000000000..0f6e19f38b28 --- /dev/null +++ b/core/services/jobs/dispatch_loop_test.go @@ -0,0 +1,177 @@ +package jobs + +import ( + "context" + "encoding/json" + "errors" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +var _ = Describe("The dispatch loop", func() { + var ( + ctx context.Context + db *gorm.DB + bus *testutil.FakeBus + picker *fakePicker + control *fakeControl + broadcast *fakeBroadcast + store *fakeStore + ) + + newLoop := func(owner string) *DispatchLoop { + GinkgoHelper() + live(db, owner) + loop, err := NewDispatchLoop(DispatchConfig{ + DB: db, Owner: owner, Picker: picker, Control: control, Broadcast: broadcast, Store: store, Hints: bus, + Interval: 50 * time.Millisecond, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(loop.Start(ctx)).To(Succeed()) + DeferCleanup(loop.Stop) + return loop + } + + rows := func() int64 { + var n int64 + Expect(db.Model(&WorkClaim{}).Count(&n).Error).To(Succeed()) + return n + } + + ci := func(jobID string) JobEvent { return JobEvent{JobID: jobID, TaskID: "t1"} } + + BeforeEach(func() { + c, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + ctx = c + db = newClaimDB() + bus = testutil.NewFakeBus() + picker = &fakePicker{nodes: []string{"w1"}, types: map[string]string{"w1": "agent"}} + control = &fakeControl{do: func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, nil + }} + broadcast = &fakeBroadcast{} + store = &fakeStore{} + }) + + It("takes a unit of each kind from the queue, drives it on a worker and removes the claim", func() { + control.do = func(_ context.Context, _, verb string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + if verb == workerctl.VerbMCPCIRun { + return workerctl.RunReply{JobID: "j1", Status: "completed", Result: "ok"}, nil + } + return workerctl.RunReply{}, nil + } + newLoop("replica-a") + queue := NewClaimQueue(db, bus) + Expect(queue.Enqueue(ctx, messaging.WorkMCPCI, ci("j1"))).To(Succeed()) + Expect(queue.Enqueue(ctx, messaging.WorkAgentRun, map[string]string{"agent_name": "a"})).To(Succeed()) + Expect(queue.Enqueue(ctx, messaging.WorkTask, ci("j2"))).To(Succeed()) + + Eventually(rows, 10*time.Second).Should(BeZero()) + verbs := map[string]bool{} + for _, c := range control.seen() { + verbs[c.verb] = true + } + Expect(verbs).To(Equal(map[string]bool{workerctl.VerbMCPCIRun: true, workerctl.VerbAgentExecute: true})) + Expect(store.terminals()).To(Equal([]terminal{{"j1", "completed", "ok", ""}})) + }) + + It("releases the claim, with a wait, when the worker could not be reached", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, errors.New("the tunnel broke") + } + newLoop("replica-a") + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, ci("j1"))).To(Succeed()) + + Eventually(func() int { + var row WorkClaim + if err := db.First(&row).Error; err != nil { + return -1 + } + return row.Attempts + }, 10*time.Second).Should(Equal(1)) + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed()) + Expect(row.State).To(Equal(ClaimPending)) + Expect(row.NotBefore).ToNot(BeNil()) + Expect(store.terminals()).To(BeEmpty()) + }) + + It("leaves the claim with its replica when the job cannot be written, and does not run the work a second time", func() { + store.fail = errors.New("database is away") + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{JobID: "j1", Status: "completed"}, nil + } + newLoop("replica-a") + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, ci("j1"))).To(Succeed()) + + Eventually(func() int { return len(control.seen()) }, 10*time.Second).Should(Equal(1)) + Consistently(func() int { return len(control.seen()) }, time.Second).Should(Equal(1)) + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed()) + Expect(row.State).To(Equal(ClaimClaimed)) + Expect(row.ClaimedBy).To(Equal("replica-a")) + }) + + It("ends the run of a job that is cancelled, closes the job as cancelled and removes the claim", func() { + started := make(chan struct{}) + control.do = func(ctx context.Context, _, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + close(started) + <-ctx.Done() + return workerctl.RunReply{}, ctx.Err() + } + newLoop("replica-a") + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, ci("j1"))).To(Succeed()) + Eventually(started, 10*time.Second).Should(BeClosed()) + + Expect(bus.Publish(messaging.SubjectJobCancel("j1"), CancelEvent{JobID: "j1"})).To(Succeed()) + + Eventually(rows, 10*time.Second).Should(BeZero()) + Expect(store.terminals()).To(Equal([]terminal{{"j1", "cancelled", "", "cancelled"}})) + }) + + It("is built only with what it needs", func() { + good := DispatchConfig{DB: db, Owner: "replica-a", Picker: picker, Control: control, Hints: bus} + for name, mutate := range map[string]func(*DispatchConfig){ + "no database": func(c *DispatchConfig) { c.DB = nil }, + "no owner": func(c *DispatchConfig) { c.Owner = "" }, + "no picker": func(c *DispatchConfig) { c.Picker = nil }, + "no control": func(c *DispatchConfig) { c.Control = nil }, + "no bus": func(c *DispatchConfig) { c.Hints = nil }, + } { + cfg := good + mutate(&cfg) + _, err := NewDispatchLoop(cfg) + Expect(err).To(HaveOccurred(), name) + } + _, err := NewDispatchLoop(good) + Expect(err).ToNot(HaveOccurred()) + }) + + It("stops claiming on Stop, and waits for the runs in flight to be settled", func() { + started := make(chan struct{}) + release := make(chan struct{}) + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + close(started) + <-release + return workerctl.RunReply{}, nil + } + loop := newLoop("replica-a") + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkAgentRun, "x")).To(Succeed()) + Eventually(started, 10*time.Second).Should(BeClosed()) + + stopped := make(chan struct{}) + go func() { defer close(stopped); loop.Stop() }() + Consistently(stopped, 200*time.Millisecond).ShouldNot(BeClosed()) + close(release) + Eventually(stopped, 5*time.Second).Should(BeClosed()) + Expect(rows()).To(BeZero()) + }) +}) diff --git a/core/services/jobs/dispatcher_test.go b/core/services/jobs/dispatcher_test.go index dbf27c84658c..ccf0c1720ec6 100644 --- a/core/services/jobs/dispatcher_test.go +++ b/core/services/jobs/dispatcher_test.go @@ -452,6 +452,23 @@ var _ = Describe("Dispatcher", func() { Expect(bus.QueueGroups()).To(BeEmpty()) }) + // Dispatcher.Cancel publishes jobs..cancel, and on the NATS carrier + // nothing subscribes to it: no worker and no frontend. A cancel of a + // distributed job therefore ends nothing there, as it has never. The + // carrier of the claim queue gives the cancel a consumer (AgentDriver), + // and these specs pin the NATS side so that a change to it is deliberate. + It("publishes the cancel of a job and subscribes to none", func() { + // What the dispatcher listens to, read before the spec adds its own. + for _, subject := range bus.SubscribedSubjects() { + Expect(messaging.SubjectMatches(subject, messaging.SubjectJobCancel("j1"))).To(BeFalse(), subject) + } + var heard []byte + _, err := bus.Subscribe(messaging.SubjectJobCancel("j1"), func(b []byte) { heard = b }) + Expect(err).ToNot(HaveOccurred()) + Expect(disp.Cancel("j1")).To(Succeed()) + Expect(heard).To(MatchJSON(`{"job_id":"j1"}`)) + }) + It("persists the result a worker publishes", func() { task := &TaskRecord{UserID: "user-1", Name: "result-task", Model: "m", Enabled: true} Expect(store.CreateTask(task)).To(Succeed()) diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index c646fa9acd62..e4139739eb94 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -109,6 +109,7 @@ const ( const ( SubjectJobResultWildcard = "jobs.*.result" SubjectJobProgressWildcard = "jobs.*.progress" + SubjectJobCancelWildcard = "jobs.*.cancel" SubjectAgentEventsWildcard = "agent.*.events.*" SubjectAgentCancelWildcard = "agent.*.cancel" SubjectGalleryCancelWildcard = "gallery.*.cancel" diff --git a/core/services/testutil/fakebus.go b/core/services/testutil/fakebus.go index 607a57bd2d1e..e59f288c20f5 100644 --- a/core/services/testutil/fakebus.go +++ b/core/services/testutil/fakebus.go @@ -83,6 +83,18 @@ func (b *FakeBus) Publish(subject string, data any) error { return nil } +// SubscribedSubjects returns the subject of each live subscription, so a spec +// can pin that a component listens to nothing on a subject. +func (b *FakeBus) SubscribedSubjects() []string { + b.mu.Lock() + defer b.mu.Unlock() + out := make([]string, 0, len(b.subs)) + for _, s := range b.subs { + out = append(out, s.subject) + } + return out +} + // PublishCount returns how many messages were published on the exact subject. func (b *FakeBus) PublishCount(subject string) int { b.mu.Lock() From 5a7473a8f1143f8ff0d2ff53a68d16c9309d4ba4 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:23:35 +0000 Subject: [PATCH 083/162] feat(nodes): demote a node whose tunnel is gone while the tunnel is the active carrier The heartbeat of a worker is an HTTP call that does not use its tunnel. A worker can send it for ever with no tunnel, and the health monitor saw a node that was up. A scheduler that marked such a node unhealthy after a failed request then watched the monitor promote it again at the next tick, and the node flapped between the two. While the tunnel is the active carrier, the monitor now reads the presence of each node that heartbeats. Only a departure older than the reconnect grace counts. A worker that is dialling again, one that never dialled and a read that failed are left alone, because none of them is evidence about the worker. A node whose tunnel is gone is demoted with a status-only change, never taken offline, so the read cannot delete rows. It is neither promoted nor probed until the tunnel is back. Agent workers are included, because they are reached through their tunnels and nothing else. On NATS nothing is read. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 7 + core/services/cluster/presence.go | 8 + core/services/nodes/health.go | 79 +++++++ core/services/nodes/health_presence_test.go | 217 ++++++++++++++++++++ 4 files changed, 311 insertions(+) create mode 100644 core/services/nodes/health_presence_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index b2c675fdd621..3cac082cfac3 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -603,6 +603,13 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, err } clusterReg := cluster.NewRegistry(authDB) + // While the tunnel is the active carrier, a node that heartbeats and holds no + // tunnel is demoted once its departure is older than the grace. On NATS no + // node holds a tunnel, and the monitor reads nothing. + healthMon.UsePresence(clusterReg, cluster.DefaultReconnectGrace, func() bool { + set := active.Load() + return set != nil && set.Name == cluster.CarrierTunnel + }) tunnels := tunnel.NewRegistry(clusterReg, cfg.Distributed.InstanceID) // If a peer sweeps this replica while it stalls, the claims of the tunnels // that it still holds are written again when it registers again. diff --git a/core/services/cluster/presence.go b/core/services/cluster/presence.go index b0b8fdf0c379..129e97263692 100644 --- a/core/services/cluster/presence.go +++ b/core/services/cluster/presence.go @@ -6,6 +6,14 @@ import ( "time" ) +// DefaultReconnectGrace is how long a departure must be older than before a +// reader may call the worker gone. A worker that moves between replicas, or +// whose network blinks, dials again within this time, and calling it gone +// sooner would demote a worker that is on its way back. It is below +// DepartedRetention, so the sweep keeps a departure for as long as the reader +// needs it. +const DefaultReconnectGrace = time.Minute + // Presence is what this package can say about a worker's tunnel. // // Four values and not a boolean, because the conditions underneath are four and diff --git a/core/services/nodes/health.go b/core/services/nodes/health.go index ad9f235f0cad..552a72f477ef 100644 --- a/core/services/nodes/health.go +++ b/core/services/nodes/health.go @@ -9,6 +9,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/advisorylock" + "github.com/mudler/LocalAI/core/services/cluster" grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/xlog" "gorm.io/gorm" @@ -32,6 +33,12 @@ type modelKey struct { ReplicaIndex int } +// NodePresenceReader says what a deployment knows about the tunnel of a worker. +// Satisfied by *cluster.Registry. +type NodePresenceReader interface { + Presence(ctx context.Context, nodeID string, grace time.Duration) (cluster.Presence, error) +} + // HealthMonitor periodically checks the health of registered backend nodes. type HealthMonitor struct { registry NodeHealthStore @@ -45,6 +52,51 @@ type HealthMonitor struct { misses map[modelKey]int // consecutive failed-probe counts; reset on success or model removal cancel context.CancelFunc cancelMu sync.Mutex + + // presence and tunnelActive are set by UsePresence. They are read by the + // loop and written before it starts. + presence NodePresenceReader + reconnectGrace time.Duration + tunnelActive func() bool +} + +// UsePresence makes the monitor read the tunnel of each node while the tunnel is +// the active carrier. A node that heartbeats over HTTP and holds no tunnel is up +// and cannot be reached, and nothing else shows it: the heartbeat is fine, and +// the verbs that fail with no route are a fact of one request. Without this +// read, a scheduler that marks such a node unhealthy on a failed request sees +// the monitor promote it again at the next tick, and the node flaps. +// +// tunnelActive is asked at each tick. When NATS is the active carrier no node +// holds a tunnel and nothing is read. A grace of zero or less takes +// cluster.DefaultReconnectGrace. Call it before Start. +func (hm *HealthMonitor) UsePresence(reader NodePresenceReader, grace time.Duration, tunnelActive func() bool) { + hm.presence = reader + hm.reconnectGrace = cmp.Or(grace, cluster.DefaultReconnectGrace) + hm.tunnelActive = tunnelActive +} + +// tunnelDeparted reports whether this deployment has decided that the tunnel of a +// node is gone: no live replica holds it, and the departure is older than the +// reconnect grace. +// +// Only cluster.PresenceGone answers true. Reconnecting is a worker that is +// dialling again, unknown is a worker that never dialled or whose departure aged +// out, and a read that fails is no answer at all. Acting on any of them would +// demote a fleet for a reason that has nothing to do with a worker. It applies to +// every node type, because an agent worker is reached through its tunnel and +// nothing else. +func (hm *HealthMonitor) tunnelDeparted(ctx context.Context, node *BackendNode) bool { + if hm.presence == nil || hm.tunnelActive == nil || !hm.tunnelActive() || node == nil { + return false + } + p, err := hm.presence.Presence(ctx, node.ID, hm.reconnectGrace) + if err != nil { + xlog.Warn("Health monitor could not read node presence; leaving the node's status alone", + "node", node.Name, "nodeID", node.ID, "error", err) + return false + } + return p == cluster.PresenceGone } // NewHealthMonitor creates a new HealthMonitor. @@ -171,6 +223,33 @@ func (hm *HealthMonitor) doCheckAll(ctx context.Context) { continue } + // The heartbeat is fresh, so the supervisor of the worker is alive. That is + // not the same as the deployment being able to reach it: the heartbeat is + // an HTTP call that does not use the tunnel, and a worker can send it for + // ever with no tunnel (a proxy that stopped upgrading websockets, a + // credential that was rotated, a reconnect loop longer than the grace). + // + // The node is demoted and not marked offline. MarkUnhealthy is status + // only, and status is enough: routing selects on status=healthy, so the + // loaded rows stop being chosen and the model is placed somewhere that can + // be reached. MarkOffline deletes the rows of the node, and deleting rows + // on a presence read would give any defect in that read the largest blast + // radius in the system. + // + // No re-promotion and no per-model probes follow: the probes would dial a + // worker that has no route and count none of it, and the promotion below + // would undo the demotion at the next tick. + if hm.tunnelDeparted(ctx, &node) { + if node.Status != StatusUnhealthy && node.Status != StatusOffline { + xlog.Warn("Node is heartbeating but its tunnel has been gone longer than the reconnect grace; marking unhealthy", + "node", node.Name, "nodeID", node.ID, "type", node.NodeType, "grace", hm.reconnectGrace) + if err := hm.registry.MarkUnhealthy(ctx, node.ID); err != nil { + xlog.Error("Failed to mark a departed node unhealthy", "node", node.Name, "error", err) + } + } + continue + } + // Heartbeat is fresh — node is alive if node.Status == StatusUnhealthy || node.Status == StatusOffline { xlog.Info("Node recovered", "node", node.Name) diff --git a/core/services/nodes/health_presence_test.go b/core/services/nodes/health_presence_test.go new file mode 100644 index 000000000000..60516bc10ded --- /dev/null +++ b/core/services/nodes/health_presence_test.go @@ -0,0 +1,217 @@ +package nodes + +import ( + "context" + "errors" + "runtime" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// fakePresence says what each node's tunnel is, and counts the questions. +type fakePresence struct { + mu sync.Mutex + by map[string]cluster.Presence + err error + graces []time.Duration +} + +func (f *fakePresence) Presence(_ context.Context, nodeID string, grace time.Duration) (cluster.Presence, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.graces = append(f.graces, grace) + if f.err != nil { + return cluster.PresenceUnknown, f.err + } + return f.by[nodeID], nil +} + +var _ = Describe("The health monitor when the tunnel is the active carrier", func() { + const ( + staleThreshold = 30 * time.Second + grace = 45 * time.Second + ) + var ( + store *fakeNodeHealthStore + factory *fakeBackendClientFactory + presence *fakePresence + active bool + hm *HealthMonitor + ) + + BeforeEach(func() { + store = newFakeNodeHealthStore() + factory = newFakeBackendClientFactory() + presence = &fakePresence{by: map[string]cluster.Presence{}} + active = true + hm = newTestHealthMonitor(store, factory, true, staleThreshold) + hm.perModelHealthCheck = true + hm.UsePresence(presence, grace, func() bool { return active }) + }) + + heartbeating := func(id, status string) { + store.addNode(makeTestNode(id, id, "", status, freshTime())) + store.addNodeModel(id, NodeModel{NodeID: id, ModelName: "m", Address: "127.0.0.1:50053"}) + } + + It("demotes a node that heartbeats and whose tunnel has been gone longer than the grace, and reaps nothing", func() { + heartbeating("n1", StatusHealthy) + presence.by["n1"] = cluster.PresenceGone + + hm.doCheckAll(context.Background()) + + Expect(store.getNode("n1").Status).To(Equal(StatusUnhealthy)) + calls := store.getCalls() + Expect(calls).To(ContainElement("MarkUnhealthy:n1")) + Expect(calls).ToNot(ContainElement(ContainSubstring("MarkOffline")), "status only: a read of presence must not delete rows") + Expect(calls).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) + Expect(presence.graces).To(Equal([]time.Duration{grace})) + }) + + It("does not probe the backends of a node it cannot reach, so no probe can count against a model", func() { + heartbeating("n1", StatusHealthy) + presence.by["n1"] = cluster.PresenceGone + factory.setClient("127.0.0.1:50053", &fakeBackendClient{healthy: false}) + for range perModelMissThreshold + 1 { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) + Expect(hm.misses).To(BeEmpty()) + }) + + DescribeTable("leaves a node alone for every answer that is not a departure older than the grace", + func(p cluster.Presence) { + heartbeating("n1", StatusHealthy) + presence.by["n1"] = p + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("MarkUnhealthy"))) + }, + Entry("a tunnel that a live replica holds", cluster.PresenceConnected), + Entry("a worker that is reconnecting now", cluster.PresenceReconnecting), + Entry("a worker that never dialled, or whose departure aged out", cluster.PresenceUnknown), + ) + + It("leaves a node alone when presence cannot be read, because a failed read is not an answer", func() { + heartbeating("n1", StatusHealthy) + presence.err = errors.New("database is away") + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("MarkUnhealthy"))) + }) + + It("does not ask about presence when NATS is the active carrier, where no node holds a tunnel", func() { + active = false + heartbeating("n1", StatusHealthy) + presence.by["n1"] = cluster.PresenceGone + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + Expect(presence.graces).To(BeEmpty()) + }) + + It("does not mark a node unhealthy again on every tick, and does not promote it while its tunnel is gone", func() { + heartbeating("n1", StatusHealthy) + presence.by["n1"] = cluster.PresenceGone + for range 4 { + hm.doCheckAll(context.Background()) + } + count := 0 + for _, c := range store.getCalls() { + if c == "MarkUnhealthy:n1" { + count++ + } + Expect(c).ToNot(Equal("MarkHealthy:n1")) + } + Expect(count).To(Equal(1)) + Expect(store.getNode("n1").Status).To(Equal(StatusUnhealthy)) + }) + + It("promotes the node again when its tunnel is back", func() { + heartbeating("n1", StatusUnhealthy) + presence.by["n1"] = cluster.PresenceConnected + hm.doCheckAll(context.Background()) + Expect(store.getCalls()).To(ContainElement("MarkHealthy:n1")) + }) + + It("reads the departure of an agent worker as it reads that of a backend worker, because an agent worker holds a tunnel too", func() { + agent := makeTestNode("a1", "a1", "", StatusHealthy, freshTime()) + agent.NodeType = NodeTypeAgent + store.addNode(agent) + presence.by["a1"] = cluster.PresenceGone + hm.doCheckAll(context.Background()) + Expect(store.getNode("a1").Status).To(Equal(StatusUnhealthy)) + }) + + It("still takes a stale heartbeat offline, as before", func() { + store.addNode(makeTestNode("n2", "n2", "", StatusHealthy, staleTime(staleThreshold))) + hm.doCheckAll(context.Background()) + Expect(store.getCalls()).To(ContainElement("MarkOffline:n2")) + }) + + It("leaves a draining node to the operator", func() { + store.addNode(makeTestNode("n3", "n3", "", StatusDraining, freshTime())) + presence.by["n3"] = cluster.PresenceGone + hm.doCheckAll(context.Background()) + Expect(store.getNode("n3").Status).To(Equal(StatusDraining)) + }) + + It("takes the default grace of the deployment when it is built with none", func() { + other := newTestHealthMonitor(store, factory, true, staleThreshold) + other.UsePresence(presence, 0, func() bool { return true }) + heartbeating("n1", StatusHealthy) + presence.by["n1"] = cluster.PresenceGone + other.doCheckAll(context.Background()) + Expect(presence.graces).ToNot(BeEmpty()) + Expect(presence.graces[0]).To(BeNumerically(">", 0)) + }) +}) + +var _ = Describe("The health monitor and the connection rows of a real cluster", func() { + It("demotes a heartbeating node whose tunnel left longer ago than the grace, and promotes it when the tunnel is held again", func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg, err := NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + clusterR := cluster.NewRegistry(db) + + node := &BackendNode{Name: "w1", NodeType: NodeTypeBackend} + Expect(reg.Register(ctx, node, true)).To(Succeed()) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + + hm := NewHealthMonitor(reg, nil, 15*time.Second, time.Hour, "", false) + hm.UsePresence(clusterR, time.Minute, func() bool { return true }) + + epoch, err := clusterR.Claim(ctx, node.ID, "replica-a") + Expect(err).ToNot(HaveOccurred()) + hm.doCheckAll(ctx) + got, err := reg.Get(ctx, node.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Status).To(Equal(StatusHealthy)) + + Expect(clusterR.Release(ctx, node.ID, "replica-a", epoch)).To(Succeed()) + hm.doCheckAll(ctx) + got, _ = reg.Get(ctx, node.ID) + Expect(got.Status).To(Equal(StatusHealthy), "inside the grace the worker is reconnecting, and nobody acts on that") + + Expect(db.Exec(`UPDATE node_connections SET disconnected_at = now() - interval '10 minutes' WHERE node_id = ?`, node.ID).Error).To(Succeed()) + hm.doCheckAll(ctx) + got, _ = reg.Get(ctx, node.ID) + Expect(got.Status).To(Equal(StatusUnhealthy)) + + _, err = clusterR.Claim(ctx, node.ID, "replica-a") + Expect(err).ToNot(HaveOccurred()) + hm.doCheckAll(ctx) + got, _ = reg.Get(ctx, node.ID) + Expect(got.Status).To(Equal(StatusHealthy)) + }) +}) From 5d8dd9fd806c2e70d9a653084122baeff54e1dfe Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:26:17 +0000 Subject: [PATCH 084/162] feat(nodes): stage files through shared object storage on either carrier The stager over shared object storage sent its worker requests as NATS request and reply, so a deployment on the tunnel could not use it and fell back to pushing every file over the bulk lane. The file verbs already exist on the control plane of the worker. S3FileStager now sends them over a control link, the abstraction the other verbs already use, so the same stager runs over NATS and over the HTTP control plane of the tunnel. NewS3NATSFileStager is the NATS form and NewS3TunnelFileStager the tunnel form. The set of the tunnel carrier takes S3Staging and a file manager, as the NATS set does. A node that no route reaches is ErrNoRoute on both, and a refusal of the worker is not. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/nats_test.go | 2 +- core/services/carrier/tunnel.go | 38 ++- core/services/nodes/control_nats.go | 10 + core/services/nodes/file_stager.go | 8 +- core/services/nodes/file_stager_s3.go | 85 ++++--- .../nodes/file_stager_s3_carriers_test.go | 234 ++++++++++++++++++ 6 files changed, 329 insertions(+), 48 deletions(-) create mode 100644 core/services/nodes/file_stager_s3_carriers_test.go diff --git a/core/services/carrier/nats_test.go b/core/services/carrier/nats_test.go index b8b285677383..0cc1b5b16cfe 100644 --- a/core/services/carrier/nats_test.go +++ b/core/services/carrier/nats_test.go @@ -67,7 +67,7 @@ var _ = Describe("NewNATSSet", func() { o.S3Staging = true set, err = carrier.NewNATSSet(o) Expect(err).ToNot(HaveOccurred()) - Expect(set.Files).To(BeAssignableToTypeOf(&nodes.S3NATSFileStager{})) + Expect(set.Files).To(BeAssignableToTypeOf(&nodes.S3FileStager{})) }) It("forwards the reconnect hook of the connection", func() { diff --git a/core/services/carrier/tunnel.go b/core/services/carrier/tunnel.go index 7fc0e04fcebb..838643b70ee0 100644 --- a/core/services/carrier/tunnel.go +++ b/core/services/carrier/tunnel.go @@ -11,6 +11,7 @@ import ( "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/storage" "github.com/mudler/LocalAI/core/services/tunnel" ) @@ -42,6 +43,13 @@ type TunnelOptions struct { // Token is the registration token. A worker checks it on its control plane, // its file routes and its backend processes. Token string + + // S3Staging selects the stager over shared object storage: the frontend puts + // a file in the store and sends the worker the file verb, over the control + // plane of the tunnel. Without it the files go over the bulk lane directly. + // FileManager is the shared store, and it is required with S3Staging. + S3Staging bool + FileManager *storage.FileManager } // NewTunnelSet builds the set of seam implementations that reach workers through @@ -61,11 +69,25 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { bulkDial := httpDialerFor(o.Dialer, tunnel.WithBulkLane()) control := nodes.NewControlClient(httpDial, o.Token) - files := nodes.NewHTTPFileStager(func(nodeID string) (string, error) { - // The host is never dialled: the dialer opens a stream on the tunnel of - // the node. It names the node in a log line and in the URL. - return nodes.WorkerHTTPHost(nodeID, ""), nil - }, o.Token, bulkDial) + var ( + files FileCarrier + forgetFile func(string) + ) + if o.S3Staging { + if o.FileManager == nil { + return nil, errors.New("tunnel set needs a file manager for S3 staging") + } + // The bytes go through the object store, and only the verbs go through + // the tunnel, so the lane of the files is not needed. + files = nodes.NewS3TunnelFileStager(o.FileManager, control) + } else { + stager := nodes.NewHTTPFileStager(func(nodeID string) (string, error) { + // The host is never dialled: the dialer opens a stream on the tunnel of + // the node. It names the node in a log line and in the URL. + return nodes.WorkerHTTPHost(nodeID, ""), nil + }, o.Token, bulkDial) + files, forgetFile = stager, stager.ForgetNode + } set := &Set{ Name: cluster.CarrierTunnel, @@ -78,10 +100,12 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { Clients: nodes.NewDialerClientFactory(o.Token, grpcDialerFor(o.Dialer)), Dialer: httpDial, Agents: o.Agents, - // The control client and the stager keep one client for each node. + // The control client and the HTTP stager keep one client for each node. ForgetNode: func(nodeID string) { control.ForgetNode(nodeID) - files.ForgetNode(nodeID) + if forgetFile != nil { + forgetFile(nodeID) + } }, Close: o.Fanout.Close, } diff --git a/core/services/nodes/control_nats.go b/core/services/nodes/control_nats.go index d2c8dc324f8c..eb79dfd99181 100644 --- a/core/services/nodes/control_nats.go +++ b/core/services/nodes/control_nats.go @@ -68,6 +68,16 @@ func natsSubject(nodeID, verb string) string { return messaging.SubjectNodeModelOp(nodeID) case workerctl.VerbModelDelete: return messaging.SubjectNodeModelDelete(nodeID) + case workerctl.VerbFilesEnsure: + return messaging.SubjectNodeFilesEnsure(nodeID) + case workerctl.VerbFilesStage: + return messaging.SubjectNodeFilesStage(nodeID) + case workerctl.VerbFilesTemp: + return messaging.SubjectNodeFilesTemp(nodeID) + case workerctl.VerbFilesListDir: + return messaging.SubjectNodeFilesListDir(nodeID) + case workerctl.VerbFilesRelease: + return messaging.SubjectNodeFilesRelease(nodeID) } return "" } diff --git a/core/services/nodes/file_stager.go b/core/services/nodes/file_stager.go index 71387426f424..805088b7212c 100644 --- a/core/services/nodes/file_stager.go +++ b/core/services/nodes/file_stager.go @@ -10,13 +10,15 @@ import ( // FileStager abstracts file transfer between frontend and backend nodes // in distributed mode. Two implementations exist: // -// 1. S3NATSFileStager (primary): Both sides have FileManager with same S3. -// Frontend uploads to S3, sends NATS request-reply to backend to download locally. +// 1. S3FileStager (primary): Both sides have FileManager with same S3. +// Frontend uploads to S3, sends the file verb to the backend (a NATS +// request-reply, or an HTTP request to its control plane on the tunnel) +// to download locally. // // 2. HTTPFileStager (fallback): Frontend pushes/pulls files directly over // HTTP to a small file transfer server on the backend node (no S3 needed). // -// S3NATSFileStager returns ErrNoRoute when nothing is listening for the node; +// S3FileStager returns ErrNoRoute when no route to the node exists; // HTTPFileStager reports connection failures as ordinary errors. See ErrNoRoute // for what a caller may do with it. type FileStager interface { diff --git a/core/services/nodes/file_stager_s3.go b/core/services/nodes/file_stager_s3.go index 5f1c2bd55e28..8ad6e478ae16 100644 --- a/core/services/nodes/file_stager_s3.go +++ b/core/services/nodes/file_stager_s3.go @@ -12,26 +12,47 @@ import ( "github.com/mudler/xlog" ) -// S3NATSFileStager implements FileStager using S3 for storage and NATS -// request-reply for coordination with backend nodes. Both frontend and +// S3FileStager implements FileStager using S3 for storage and the file verbs of +// the control plane for coordination with backend nodes. Both frontend and // backend nodes share the same S3 bucket. The flow is: // // 1. Frontend uploads file to S3 -// 2. Frontend sends NATS request to nodes.{nodeID}.files.ensure +// 2. Frontend sends the files.ensure verb to the node // 3. Backend downloads from S3 to local cache, replies with local path -type S3NATSFileStager struct { +// +// How the verbs travel is the link: a NATS request to the subject of the node, or +// an HTTP request to the control plane of the worker through its tunnel. The +// stager is the same for both, so a deployment with shared object storage stages +// its files the same way on either carrier. The link maps a node that nothing can +// reach onto ErrNoRoute. +type S3FileStager struct { fm *storage.FileManager - nats messaging.MessagingClient + link controlLink +} + +// NewS3NATSFileStager creates a file stager that sends the file verbs over NATS. +func NewS3NATSFileStager(fm *storage.FileManager, nats messaging.MessagingClient) *S3FileStager { + return &S3FileStager{fm: fm, link: &natsLink{bus: nats}} } -// NewS3NATSFileStager creates a new S3+NATS file stager. -func NewS3NATSFileStager(fm *storage.FileManager, nats messaging.MessagingClient) *S3NATSFileStager { - return &S3NATSFileStager{fm: fm, nats: nats} +// NewS3TunnelFileStager creates a file stager that sends the file verbs to the +// HTTP control plane of a worker, through whatever dialer client has. +func NewS3TunnelFileStager(fm *storage.FileManager, client *ControlClient) *S3FileStager { + return &S3FileStager{fm: fm, link: &httpLink{client: client}} +} + +// fileVerb sends one file verb to a node and decodes the reply. +func fileVerb[Req, Reply any](ctx context.Context, s *S3FileStager, nodeID, verb string, req Req, timeout time.Duration) (*Reply, error) { + var reply Reply + if err := s.link.request(ctx, nodeID, verb, req, &reply, timeout); err != nil { + return nil, err + } + return &reply, nil } // EnsureRemote uploads a local file to S3 (if not already there) and sends // a NATS request-reply to the backend node to download it locally. -func (s *S3NATSFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) { +func (s *S3FileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) { // Upload to S3 if not already present exists, _ := s.fm.Exists(ctx, key) if !exists { @@ -47,9 +68,8 @@ func (s *S3NATSFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, } } - // Send NATS request-reply to backend - subject := messaging.SubjectNodeFilesEnsure(nodeID) - reply, err := controlRequestJSON[workerctl.FileEnsureRequest, workerctl.FileEnsureReply](s.nats, subject, workerctl.FileEnsureRequest{Key: key}, 10*time.Minute) + // Ask the backend to download it + reply, err := fileVerb[workerctl.FileEnsureRequest, workerctl.FileEnsureReply](ctx, s, nodeID, workerctl.VerbFilesEnsure, workerctl.FileEnsureRequest{Key: key}, 10*time.Minute) if err != nil { return "", err } @@ -62,7 +82,7 @@ func (s *S3NATSFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, } // FetchRemote tells the backend to upload a file to S3, then downloads it locally. -func (s *S3NATSFileStager) FetchRemote(ctx context.Context, nodeID, remotePath, localDst string) error { +func (s *S3FileStager) FetchRemote(ctx context.Context, nodeID, remotePath, localDst string) error { // Tell backend to upload to S3 key := storage.EphemeralKey(remotePath, "fetch", "output") return s.fetchRemoteWithKey(ctx, nodeID, remotePath, key, localDst, true) @@ -70,16 +90,15 @@ func (s *S3NATSFileStager) FetchRemote(ctx context.Context, nodeID, remotePath, // FetchRemoteByKey tells the backend to upload a file (identified by key) to S3, // then downloads it locally. The key is used as-is for S3 routing. -func (s *S3NATSFileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, localDst string) error { +func (s *S3FileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, localDst string) error { // For S3 mode, we still need the remote path — derive it from the key. // The backend serves the file from its data dir based on the key prefix. remotePath := "/" + key // e.g. "/data/quantization/{jobID}/model.gguf" return s.fetchRemoteWithKey(ctx, nodeID, remotePath, key, localDst, true) } -func (s *S3NATSFileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePath, key, localDst string, cleanup bool) error { - subject := messaging.SubjectNodeFilesStage(nodeID) - reply, err := controlRequestJSON[workerctl.FileStageRequest, workerctl.FileStageReply](s.nats, subject, workerctl.FileStageRequest{LocalPath: remotePath, Key: key}, 10*time.Minute) +func (s *S3FileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remotePath, key, localDst string, cleanup bool) error { + reply, err := fileVerb[workerctl.FileStageRequest, workerctl.FileStageReply](ctx, s, nodeID, workerctl.VerbFilesStage, workerctl.FileStageRequest{LocalPath: remotePath, Key: key}, 10*time.Minute) if err != nil { return err } @@ -106,10 +125,9 @@ func (s *S3NATSFileStager) fetchRemoteWithKey(ctx context.Context, nodeID, remot return nil } -// AllocRemoteTemp asks the backend to allocate a temp file via NATS request-reply. -func (s *S3NATSFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) { - subject := messaging.SubjectNodeFilesTemp(nodeID) - reply, err := controlRequestJSON[workerctl.FileTempRequest, workerctl.FileTempReply](s.nats, subject, workerctl.FileTempRequest{}, 30*time.Second) +// AllocRemoteTemp asks the backend to allocate a temp file. +func (s *S3FileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) { + reply, err := fileVerb[workerctl.FileTempRequest, workerctl.FileTempReply](ctx, s, nodeID, workerctl.VerbFilesTemp, workerctl.FileTempRequest{}, 30*time.Second) if err != nil { return "", err } @@ -120,9 +138,8 @@ func (s *S3NATSFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) ( return reply.LocalPath, nil } -func (s *S3NATSFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) { - subject := messaging.SubjectNodeFilesListDir(nodeID) - reply, err := controlRequestJSON[workerctl.FileListDirRequest, workerctl.FileListDirReply](s.nats, subject, workerctl.FileListDirRequest{KeyPrefix: keyPrefix}, 30*time.Second) +func (s *S3FileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) { + reply, err := fileVerb[workerctl.FileListDirRequest, workerctl.FileListDirReply](ctx, s, nodeID, workerctl.VerbFilesListDir, workerctl.FileListDirRequest{KeyPrefix: keyPrefix}, 30*time.Second) if err != nil { return nil, err } @@ -134,9 +151,8 @@ func (s *S3NATSFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix } // StageRemoteToStore tells the backend to upload a local file to S3. -func (s *S3NATSFileStager) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error { - subject := messaging.SubjectNodeFilesStage(nodeID) - reply, err := controlRequestJSON[workerctl.FileStageRequest, workerctl.FileStageReply](s.nats, subject, workerctl.FileStageRequest{LocalPath: remotePath, Key: key}, 10*time.Minute) +func (s *S3FileStager) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error { + reply, err := fileVerb[workerctl.FileStageRequest, workerctl.FileStageReply](ctx, s, nodeID, workerctl.VerbFilesStage, workerctl.FileStageRequest{LocalPath: remotePath, Key: key}, 10*time.Minute) if err != nil { return err } @@ -149,7 +165,7 @@ func (s *S3NATSFileStager) StageRemoteToStore(ctx context.Context, nodeID, remot // ReleaseRemote evicts one exact ephemeral key from the worker before deleting // the shared object. -func (s *S3NATSFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error { +func (s *S3FileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error { if err := validateEphemeralReleaseKey(key); err != nil { return err } @@ -162,10 +178,10 @@ func (s *S3NATSFileStager) ReleaseRemote(ctx context.Context, nodeID, key string return nil } -// ReleaseRemoteRequest evicts one inference's inputs with one NATS round trip. +// ReleaseRemoteRequest evicts one inference's inputs with one round trip. // A worker that only understands the exact-key payload returns an error, so the // frontend retries each key during a rolling upgrade. -func (s *S3NATSFileStager) ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error { +func (s *S3FileStager) ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error { if err := validateEphemeralRequestRelease(requestID, keys); err != nil { return err } @@ -190,7 +206,7 @@ func (s *S3NATSFileStager) ReleaseRemoteRequest(ctx context.Context, nodeID, req return errors.Join(deleteErrors...) } -func (s *S3NATSFileStager) releaseWorkerKeys(ctx context.Context, nodeID string, request workerctl.FileReleaseRequest) error { +func (s *S3FileStager) releaseWorkerKeys(ctx context.Context, nodeID string, request workerctl.FileReleaseRequest) error { if err := ctx.Err(); err != nil { return err } @@ -202,12 +218,7 @@ func (s *S3NATSFileStager) releaseWorkerKeys(ctx context.Context, nodeID string, } timeout = min(timeout, remaining) } - reply, err := controlRequestJSON[workerctl.FileReleaseRequest, workerctl.FileReleaseReply]( - s.nats, - messaging.SubjectNodeFilesRelease(nodeID), - request, - timeout, - ) + reply, err := fileVerb[workerctl.FileReleaseRequest, workerctl.FileReleaseReply](ctx, s, nodeID, workerctl.VerbFilesRelease, request, timeout) if err != nil { if contextErr := ctx.Err(); contextErr != nil { return contextErr diff --git a/core/services/nodes/file_stager_s3_carriers_test.go b/core/services/nodes/file_stager_s3_carriers_test.go new file mode 100644 index 000000000000..c6b08d7c4b22 --- /dev/null +++ b/core/services/nodes/file_stager_s3_carriers_test.go @@ -0,0 +1,234 @@ +package nodes + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "strings" + "sync" + "time" + + "github.com/nats-io/nats.go" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// scriptedWorker is a worker whose answers to the file verbs a spec scripts. A +// rig makes it reachable by one carrier. +type scriptedWorker struct { + mu sync.Mutex + replies map[string]any // verb -> reply + requests map[string][]json.RawMessage +} + +func newScriptedWorker() *scriptedWorker { + return &scriptedWorker{replies: map[string]any{}, requests: map[string][]json.RawMessage{}} +} + +func (w *scriptedWorker) on(verb string, reply any) { + w.mu.Lock() + defer w.mu.Unlock() + w.replies[verb] = reply +} + +// answer records a request and returns the scripted reply as JSON. +func (w *scriptedWorker) answer(verb string, body []byte) []byte { + w.mu.Lock() + defer w.mu.Unlock() + w.requests[verb] = append(w.requests[verb], append(json.RawMessage(nil), body...)) + reply, ok := w.replies[verb] + if !ok { + reply = struct{}{} + } + raw, _ := json.Marshal(reply) + return raw +} + +func (w *scriptedWorker) seen(verb string) []json.RawMessage { + w.mu.Lock() + defer w.mu.Unlock() + return append([]json.RawMessage(nil), w.requests[verb]...) +} + +// natsRig reaches the scripted worker over NATS request and reply. +type natsRig struct { + releaseTestMessaging + worker *scriptedWorker + noRoute bool +} + +func (r *natsRig) Request(subject string, data []byte, _ time.Duration) ([]byte, error) { + if r.noRoute { + return nil, nats.ErrNoResponders + } + for _, verb := range []string{workerctl.VerbFilesEnsure, workerctl.VerbFilesStage, workerctl.VerbFilesTemp, workerctl.VerbFilesListDir, workerctl.VerbFilesRelease} { + if subject == natsSubject("n1", verb) { + return r.worker.answer(verb, data), nil + } + } + return nil, fmt.Errorf("the rig has no verb for %s", subject) +} + +var _ = Describe("The S3 file stager over each carrier of the file verbs", func() { + type rig struct { + name string + build func(fm *storage.FileManager, w *scriptedWorker) (FileStager, func(noRoute bool)) + release func(FileStager) RequestFileReleaser + } + + rigs := []rig{ + { + name: "NATS", + build: func(fm *storage.FileManager, w *scriptedWorker) (FileStager, func(bool)) { + bus := &natsRig{worker: w} + return NewS3NATSFileStager(fm, bus), func(noRoute bool) { bus.noRoute = noRoute } + }, + }, + { + name: "the HTTP control plane of the tunnel", + build: func(fm *storage.FileManager, w *scriptedWorker) (FileStager, func(bool)) { + worker := newControlWorker() + for _, verb := range []string{workerctl.VerbFilesEnsure, workerctl.VerbFilesStage, workerctl.VerbFilesTemp, workerctl.VerbFilesListDir, workerctl.VerbFilesRelease} { + worker.on(verb, func(rw http.ResponseWriter, r *http.Request) { + body := make([]byte, 0, 256) + buf := make([]byte, 256) + for { + n, err := r.Body.Read(buf) + body = append(body, buf[:n]...) + if err != nil { + break + } + } + rw.Header().Set("Content-Type", "application/json") + _, _ = rw.Write(w.answer(strings.ReplaceAll(strings.TrimPrefix(r.URL.Path, workerctl.Prefix), "/", "."), body)) + }) + } + stager := NewS3TunnelFileStager(fm, NewControlClient(worker.dialerFor(), "token")) + return stager, func(noRoute bool) { + if noRoute { + worker.failDial(func(context.Context) error { return ErrNoRoute }) + } else { + worker.failDial(nil) + } + } + }, + }, + } + + for _, r := range rigs { + Describe(r.name, func() { + var ( + ctx context.Context + store storage.ObjectStore + fm *storage.FileManager + worker *scriptedWorker + stager FileStager + setDown func(bool) + ) + + BeforeEach(func() { + ctx = context.Background() + var err error + store, err = storage.NewFilesystemStore(GinkgoT().TempDir()) + Expect(err).ToNot(HaveOccurred()) + fm, err = storage.NewFileManager(store, GinkgoT().TempDir()) + Expect(err).ToNot(HaveOccurred()) + worker = newScriptedWorker() + stager, setDown = r.build(fm, worker) + }) + + It("uploads a file and asks the worker to ensure it, and returns the path the worker gives", func() { + local := GinkgoT().TempDir() + "/in.bin" + Expect(writeFile(local, "payload")).To(Succeed()) + worker.on(workerctl.VerbFilesEnsure, workerctl.FileEnsureReply{LocalPath: "/cache/in.bin"}) + + path, err := stager.EnsureRemote(ctx, "n1", local, "ephemeral/r1/in.bin") + Expect(err).ToNot(HaveOccurred()) + Expect(path).To(Equal("/cache/in.bin")) + + exists, err := store.Exists(ctx, "ephemeral/r1/in.bin") + Expect(err).ToNot(HaveOccurred()) + Expect(exists).To(BeTrue()) + Expect(worker.seen(workerctl.VerbFilesEnsure)).To(HaveLen(1)) + Expect(string(worker.seen(workerctl.VerbFilesEnsure)[0])).To(MatchJSON(`{"key":"ephemeral/r1/in.bin"}`)) + }) + + It("reports a refusal of the worker as an error that is not a missing route", func() { + local := GinkgoT().TempDir() + "/in.bin" + Expect(writeFile(local, "payload")).To(Succeed()) + worker.on(workerctl.VerbFilesEnsure, workerctl.FileEnsureReply{Error: "disk full"}) + + _, err := stager.EnsureRemote(ctx, "n1", local, "k1") + Expect(err).To(MatchError(ContainSubstring("disk full"))) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) + }) + + It("reports a node that nothing can reach as ErrNoRoute", func() { + local := GinkgoT().TempDir() + "/in.bin" + Expect(writeFile(local, "payload")).To(Succeed()) + setDown(true) + _, err := stager.EnsureRemote(ctx, "n1", local, "k1") + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue(), "got %v", err) + _, err = stager.AllocRemoteTemp(ctx, "n1") + Expect(errors.Is(err, ErrNoRoute)).To(BeTrue(), "got %v", err) + }) + + It("allocates a temp file, lists a directory and stages a file to the store through the worker", func() { + worker.on(workerctl.VerbFilesTemp, workerctl.FileTempReply{LocalPath: "/tmp/x"}) + worker.on(workerctl.VerbFilesListDir, workerctl.FileListDirReply{Files: []string{"a", "b"}}) + path, err := stager.AllocRemoteTemp(ctx, "n1") + Expect(err).ToNot(HaveOccurred()) + Expect(path).To(Equal("/tmp/x")) + files, err := stager.ListRemoteDir(ctx, "n1", "data/") + Expect(err).ToNot(HaveOccurred()) + Expect(files).To(Equal([]string{"a", "b"})) + + Expect(stager.StageRemoteToStore(ctx, "n1", "/data/out.bin", "out/key")).To(Succeed()) + Expect(string(worker.seen(workerctl.VerbFilesStage)[0])).To(MatchJSON(`{"local_path":"/data/out.bin","key":"out/key"}`)) + }) + + It("fetches a file by having the worker stage it to the store, then reading it from there", func() { + Expect(store.Put(ctx, "out/key", strings.NewReader("result"))).To(Succeed()) + dst := GinkgoT().TempDir() + "/dst.bin" + Expect(stager.FetchRemoteByKey(ctx, "n1", "out/key", dst)).To(Succeed()) + got, err := readFile(dst) + Expect(err).ToNot(HaveOccurred()) + Expect(got).To(Equal("result")) + Expect(worker.seen(workerctl.VerbFilesStage)).To(HaveLen(1)) + }) + + It("evicts the worker's copy before it deletes the shared object, in one round trip for a request", func() { + keys := []string{"ephemeral/audio/r1/in.wav", "ephemeral/images/r1/f.jpg"} + for _, k := range keys { + Expect(store.Put(ctx, k, strings.NewReader("x"))).To(Succeed()) + } + releaser, ok := stager.(RequestFileReleaser) + Expect(ok).To(BeTrue()) + Expect(releaser.ReleaseRemoteRequest(ctx, "n1", "r1", keys)).To(Succeed()) + + Expect(worker.seen(workerctl.VerbFilesRelease)).To(HaveLen(1)) + Expect(string(worker.seen(workerctl.VerbFilesRelease)[0])).To(MatchJSON(`{"request_id":"r1"}`)) + for _, k := range keys { + exists, _ := store.Exists(ctx, k) + Expect(exists).To(BeFalse(), k) + } + }) + }) + } +}) + +var _ messaging.MessagingClient = (*natsRig)(nil) + +func writeFile(path, content string) error { return os.WriteFile(path, []byte(content), 0o600) } + +func readFile(path string) (string, error) { + b, err := os.ReadFile(path) + return string(b), err +} From 3a0100327e8c16f86f8ec9a4c29e6955d4c1fa26 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:30:55 +0000 Subject: [PATCH 085/162] test(carrier): run agent workers and the claim queue through the tunnel carrier The set of the tunnel carrier now builds its agent RPC itself, from a selector and the same control client as its other verbs, and takes the claim queue as its queue. The integration test uses both instead of fakes. The new specs put an agent worker on the tunnel of one replica and the dispatch loop on another. An MCP CI job runs through the peer link, its progress, its trace and its result are heard by a third subscriber on the other replica, and its terminal state is written. A subject that the worker may not ask for is refused. The events of an agent run, built by the real event bridge bound to the publisher of the delivery, arrive with their subjects and payloads. A cancel of the job ends the run on the worker with its request and closes the job as cancelled. The MCP requests reach the worker from both replicas, a tool that failed is an answer, and no worker is ErrNoRoute. The last spec kills the replica that drives a run: its database connection and its heartbeat stop and its streams end. The claim goes back to the queue, another replica drives it on the worker after it dials that replica, and the job ends. It ran twice and ended once. Staging through shared object storage is checked on the same carrier: the file is put in the store and only the verb goes over the tunnel. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/tunnel.go | 25 +- .../carrier/tunnel_agents_integration_test.go | 367 ++++++++++++++++++ .../carrier/tunnel_integration_test.go | 105 ++++- 3 files changed, 482 insertions(+), 15 deletions(-) create mode 100644 core/services/carrier/tunnel_agents_integration_test.go diff --git a/core/services/carrier/tunnel.go b/core/services/carrier/tunnel.go index 838643b70ee0..c122c1c70709 100644 --- a/core/services/carrier/tunnel.go +++ b/core/services/carrier/tunnel.go @@ -7,7 +7,6 @@ import ( "net" "time" - mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" @@ -17,11 +16,12 @@ import ( // TunnelOptions is what NewTunnelSet needs. // -// The set has seven members. Four of them belong to this carrier and are built -// here: the control verbs, the file stager, the client factory and the dialer of -// the HTTP server of a worker. The fan-out comes from NewPgbusFanout. The queue -// and the agent RPC of the tunnel carrier are provided by the caller, because the -// parts that implement them are separate from the parts that reach a worker. +// The set has seven members. Five of them belong to this carrier and are built +// here: the control verbs, the file stager, the client factory, the dialer of the +// HTTP server of a worker and the agent RPC, which picks an agent worker with the +// selector and sends the request over the same control client as the other verbs. +// The fan-out comes from NewPgbusFanout. The queue is provided by the caller +// (jobs.ClaimQueue), because it lives in the database and reaches no worker. type TunnelOptions struct { // Epoch is the epoch of the cluster row the set is built for. Epoch int64 @@ -32,9 +32,10 @@ type TunnelOptions struct { // Fanout carries the broadcasts. Fanout *Fanout - // WorkQueue and Agents are the other two members that the set needs. + // WorkQueue is the queue of the carrier. WorkQueue messaging.WorkQueue - Agents mcpTools.AgentControl + // AgentSelector picks the agent worker that gets an MCP request. + AgentSelector *nodes.AgentSelector // Registry finds the nodes that hold a model, for the verbs that fan out. Registry nodes.ModelLocator @@ -61,6 +62,12 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { if o.Fanout == nil { return nil, errors.New("tunnel set needs a fan-out") } + if o.WorkQueue == nil { + return nil, errors.New("tunnel set needs a work queue") + } + if o.AgentSelector == nil { + return nil, errors.New("tunnel set needs an agent selector") + } // The control verbs, the logs and the health of a worker go over the // inference lane. The file transfer goes over the bulk lane and never falls @@ -99,7 +106,7 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { Files: files, Clients: nodes.NewDialerClientFactory(o.Token, grpcDialerFor(o.Dialer)), Dialer: httpDial, - Agents: o.Agents, + Agents: nodes.NewAgentControlClient(o.AgentSelector, control), // The control client and the HTTP stager keep one client for each node. ForgetNode: func(nodeID string) { control.ForgetNode(nodeID) diff --git a/core/services/carrier/tunnel_agents_integration_test.go b/core/services/carrier/tunnel_agents_integration_test.go new file mode 100644 index 000000000000..63d7e77fb43b --- /dev/null +++ b/core/services/carrier/tunnel_agents_integration_test.go @@ -0,0 +1,367 @@ +package carrier_test + +import ( + "context" + "encoding/json" + "errors" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/logger" + + "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/agentworker" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/jobs" + mcpremote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" +) + +const agentToken = "own-tunnel-credential-of-the-agent" + +// heard keeps what a subscriber on a replica received. +type heard struct { + mu sync.Mutex + msgs [][]byte +} + +func (h *heard) handler(b []byte) { + h.mu.Lock() + defer h.mu.Unlock() + h.msgs = append(h.msgs, append([]byte(nil), b...)) +} + +func (h *heard) all() [][]byte { + h.mu.Lock() + defer h.mu.Unlock() + return append([][]byte(nil), h.msgs...) +} + +func (h *heard) count() int { return len(h.all()) } + +var _ = Describe("Agent workers on the tunnel carrier, end to end", func() { + var ( + ctx context.Context + db *gorm.DB + dsn string + clusterR *cluster.Registry + nodeReg *nodes.NodeRegistry + store *jobs.JobStore + a, b *replica + agentID string + + work *agentworker.Work + cfg agentworker.Config + ) + + // openDB opens a connection pool of its own, as the process of a replica has. + openDB := func() *gorm.DB { + GinkgoHelper() + conn, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + Expect(err).ToNot(HaveOccurred()) + return conn + } + + // startAgent connects the agent worker to a frontend. It can be called again + // for another frontend, as a worker that lost its replica does. + startAgent := func(frontend string) *agentworker.Runtime { + GinkgoHelper() + rt, err := agentworker.Start(ctx, agentworker.Options{ + FrontendURL: frontend, + NodeID: agentID, + TunnelToken: func() string { return agentToken }, + ControlToken: registrationToken, + Handler: agentworker.Handler(cfg, work), + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(rt.Close) + return rt + } + + // startLoop starts the dispatch loop of a replica over a connection. + startLoop := func(r *replica, conn *gorm.DB) (*jobs.DispatchLoop, context.CancelFunc) { + GinkgoHelper() + loopCtx, cancel := context.WithCancel(ctx) + DeferCleanup(cancel) + loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{ + DB: conn, Owner: r.id, + Picker: r.selector, + Control: nodes.NewControlClient(r.set.Dialer, registrationToken), + Broadcast: nodes.NewRebroadcaster(r.bus), + Store: store, + Hints: r.bus, + Interval: 100 * time.Millisecond, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(loop.Start(loopCtx)).To(Succeed()) + return loop, cancel + } + + newJob := func() (jobID, taskID string) { + GinkgoHelper() + task := &jobs.TaskRecord{UserID: "u1", Name: "t", Model: "m", Enabled: true} + Expect(store.CreateTask(task)).To(Succeed()) + job := &jobs.JobRecord{TaskID: task.ID, UserID: "u1", Status: "running", TriggeredBy: "manual"} + Expect(store.CreateJob(job)).To(Succeed()) + return job.ID, task.ID + } + + jobStatus := func(id string) func() string { + return func() string { + job, err := store.GetJob(id) + if err != nil { + return err.Error() + } + return job.Status + } + } + + claims := func() int64 { + var n int64 + Expect(db.Model(&jobs.WorkClaim{}).Count(&n).Error).To(Succeed()) + return n + } + + BeforeEach(func() { + ctx = context.Background() + db, dsn = testutil.SetupTestDBWithDSN() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + Expect(jobs.MigrateClaims(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + var err error + nodeReg, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + store, err = jobs.NewJobStore(db) + Expect(err).ToNot(HaveOccurred()) + + agent := &nodes.BackendNode{Name: "agent-1", NodeType: nodes.NodeTypeAgent, TokenHash: hashOf(registrationToken)} + Expect(nodeReg.Register(ctx, agent, true)).To(Succeed()) + agentID = agent.ID + Expect(nodeReg.SetTunnelTokenHash(ctx, agentID, hashOf(agentToken))).To(Succeed()) + + a, b = startReplica(ctx, "replica-a", db, dsn, clusterR, nodeReg), startReplica(ctx, "replica-b", db, dsn, clusterR, nodeReg) + + work = agentworker.NewWork() + cfg = agentworker.Config{ + MCPTool: agentworker.Unary(func(_ context.Context, req mcpremote.MCPToolRequest) mcpremote.MCPToolResponse { + if req.ToolName == "fails" { + return mcpremote.MCPToolResponse{Error: "the tool failed"} + } + return mcpremote.MCPToolResponse{Result: "ran " + req.ToolName} + }), + } + }) + + holds := func(r *replica) func() bool { + return func() bool { return r.tunnels.Holds(agentID) } + } + + Describe("an MCP CI job", func() { + It("runs on an agent worker that another replica holds, streams its events to every replica and records its result", func() { + var runs atomic.Int32 + _, err := work.Consume(ctx, messaging.WorkMCPCI, 1, func(_ context.Context, payload []byte, events messaging.Publisher) error { + runs.Add(1) + var evt jobs.JobEvent + Expect(json.Unmarshal(payload, &evt)).To(Succeed()) + jobs.PublishJobProgress(events, evt.JobID, "running", "started on the agent worker") + Expect(events.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{JobID: evt.JobID, TraceType: "content", TraceContent: "hello"})).To(Succeed()) + // A worker may ask for the events of a run and for nothing else. + Expect(events.Publish(messaging.SubjectJobCancel(evt.JobID), jobs.CancelEvent{JobID: evt.JobID})).To(Succeed()) + jobs.PublishJobResult(events, evt.JobID, "completed", "42", "") + return nil + }) + Expect(err).ToNot(HaveOccurred()) + + var progress, results, cancels heard + _, err = b.bus.Subscribe(messaging.SubjectJobProgressWildcard, progress.handler) + Expect(err).ToNot(HaveOccurred()) + _, err = b.bus.Subscribe(messaging.SubjectJobResultWildcard, results.handler) + Expect(err).ToNot(HaveOccurred()) + _, err = b.bus.Subscribe(messaging.SubjectJobCancelWildcard, cancels.handler) + Expect(err).ToNot(HaveOccurred()) + + // B holds the tunnel of the agent worker and A drives the job, so the + // run crosses the link between the replicas. + startAgent(b.url) + Eventually(holds(b), "15s", "50ms").Should(BeTrue()) + startLoop(a, openDB()) + + jobID, taskID := newJob() + Expect(a.queue.Enqueue(ctx, messaging.WorkMCPCI, jobs.JobEvent{JobID: jobID, TaskID: taskID, UserID: "u1"})).To(Succeed()) + + Eventually(jobStatus(jobID), "30s", "100ms").Should(Equal("completed")) + job, err := store.GetJob(jobID) + Expect(err).ToNot(HaveOccurred()) + Expect(job.Result).To(Equal("42")) + Eventually(claims, "10s", "100ms").Should(BeZero()) + Expect(runs.Load()).To(BeEquivalentTo(1)) + + Eventually(results.count, "10s").Should(BeNumerically(">=", 1)) + Eventually(progress.count, "10s").Should(BeNumerically(">=", 3)) + var sawTrace bool + for _, raw := range progress.all() { + var ev jobs.ProgressEvent + Expect(json.Unmarshal(raw, &ev)).To(Succeed()) + if ev.TraceContent == "hello" { + sawTrace = true + } + } + Expect(sawTrace).To(BeTrue(), "the trace of the run reached the other replica") + Consistently(cancels.count, "500ms").Should(BeZero(), "a subject the worker may not ask for is refused") + }) + + It("is closed as cancelled, and its run ends with its request, when the job is cancelled", func() { + started := make(chan struct{}) + ended := make(chan error, 1) + _, err := work.Consume(ctx, messaging.WorkMCPCI, 1, func(ctx context.Context, _ []byte, _ messaging.Publisher) error { + close(started) + <-ctx.Done() + ended <- ctx.Err() + return ctx.Err() + }) + Expect(err).ToNot(HaveOccurred()) + startAgent(a.url) + Eventually(holds(a), "15s", "50ms").Should(BeTrue()) + startLoop(a, openDB()) + + jobID, taskID := newJob() + Expect(a.queue.Enqueue(ctx, messaging.WorkMCPCI, jobs.JobEvent{JobID: jobID, TaskID: taskID, UserID: "u1"})).To(Succeed()) + Eventually(started, "20s").Should(BeClosed()) + + // The cancel is a broadcast, published by any replica. + Expect(b.bus.Publish(messaging.SubjectJobCancel(jobID), jobs.CancelEvent{JobID: jobID})).To(Succeed()) + + Eventually(ended, "20s").Should(Receive(MatchError(context.Canceled))) + Eventually(jobStatus(jobID), "20s", "100ms").Should(Equal("cancelled")) + Eventually(claims, "10s", "100ms").Should(BeZero()) + }) + }) + + Describe("an agent run", func() { + It("publishes the events of the run, built by the real event bridge, where the other replicas hear them", func() { + _, err := work.Consume(ctx, messaging.WorkAgentRun, 0, func(_ context.Context, payload []byte, events messaging.Publisher) error { + var evt agents.AgentChatEvent + Expect(json.Unmarshal(payload, &evt)).To(Succeed()) + bridge := agents.NewEventBridge(nil, nil, "agent-worker").WithPublisher(events) + Expect(bridge.PublishStatus(evt.AgentName, evt.UserID, "processing")).To(Succeed()) + Expect(bridge.PublishMessage(evt.AgentName, evt.UserID, agents.RoleAgent, "done", evt.MessageID+"-reply")).To(Succeed()) + return nil + }) + Expect(err).ToNot(HaveOccurred()) + + var seenOnA, seenOnB heard + _, err = a.bus.Subscribe(messaging.SubjectAgentEvents("helper", "alice"), seenOnA.handler) + Expect(err).ToNot(HaveOccurred()) + _, err = b.bus.Subscribe(messaging.SubjectAgentEventsWildcard, seenOnB.handler) + Expect(err).ToNot(HaveOccurred()) + + startAgent(a.url) + Eventually(holds(a), "15s", "50ms").Should(BeTrue()) + startLoop(b, openDB()) + + Expect(a.queue.Enqueue(ctx, messaging.WorkAgentRun, agents.AgentChatEvent{AgentName: "helper", UserID: "alice", Message: "hi", MessageID: "m1"})).To(Succeed()) + + Eventually(seenOnA.count, "30s", "100ms").Should(Equal(2)) + Eventually(seenOnB.count, "10s", "100ms").Should(Equal(2)) + var types []string + for _, raw := range seenOnA.all() { + var ev agents.AgentEvent + Expect(json.Unmarshal(raw, &ev)).To(Succeed()) + types = append(types, ev.EventType) + } + Expect(types).To(Equal([]string{"json_message_status", "json_message"})) + Eventually(claims, "10s", "100ms").Should(BeZero()) + }) + }) + + Describe("the MCP requests of the frontend", func() { + It("reach the agent worker from the replica that holds its tunnel and from the one that does not", func() { + startAgent(a.url) + Eventually(holds(a), "15s", "50ms").Should(BeTrue()) + for _, r := range []*replica{a, b} { + reply, err := r.agents.ExecuteMCPTool(ctx, mcpremote.MCPToolRequest{ModelName: "m", ToolName: "search"}) + Expect(err).ToNot(HaveOccurred(), r.id) + Expect(reply.Result).To(Equal("ran search"), r.id) + + failed, err := r.agents.ExecuteMCPTool(ctx, mcpremote.MCPToolRequest{ModelName: "m", ToolName: "fails"}) + Expect(err).ToNot(HaveOccurred(), "a tool that failed is an answer: "+r.id) + Expect(failed.Error).To(Equal("the tool failed"), r.id) + } + }) + + It("find no route when no agent worker is connected, which the frontend reads as no agent worker", func() { + for _, r := range []*replica{a, b} { + _, err := r.agents.ExecuteMCPTool(ctx, mcpremote.MCPToolRequest{ModelName: "m", ToolName: "search"}) + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeTrue(), "%s: %v", r.id, err) + Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse(), "absence must not reach the caller") + } + }) + }) + + Describe("a replica that dies in the middle of a run", func() { + It("loses no job: the claim goes back to the queue, another replica drives it and the job ends", func() { + var runs atomic.Int32 + started := make(chan struct{}) + _, err := work.Consume(ctx, messaging.WorkMCPCI, 0, func(ctx context.Context, payload []byte, events messaging.Publisher) error { + var evt jobs.JobEvent + Expect(json.Unmarshal(payload, &evt)).To(Succeed()) + if runs.Add(1) == 1 { + // The first run is the one whose replica dies. It ends when its + // request ends. + close(started) + <-ctx.Done() + return ctx.Err() + } + jobs.PublishJobResult(events, evt.JobID, "completed", "recovered", "") + return nil + }) + Expect(err).ToNot(HaveOccurred()) + + // A holds the tunnel of the agent worker, claims the job and drives it. + // It has a connection of its own to the database, as a process has. + first := startAgent(a.url) + Eventually(holds(a), "15s", "50ms").Should(BeTrue()) + aDB := openDB() + _, killA := startLoop(a, aDB) + + jobID, taskID := newJob() + Expect(a.queue.Enqueue(ctx, messaging.WorkMCPCI, jobs.JobEvent{JobID: jobID, TaskID: taskID, UserID: "u1"})).To(Succeed()) + Eventually(started, "20s").Should(BeClosed()) + var claim jobs.WorkClaim + Expect(db.First(&claim).Error).To(Succeed()) + Expect(claim.ClaimedBy).To(Equal("replica-a")) + Expect(claim.State).To(Equal(jobs.ClaimClaimed)) + + // A dies: its connection to the database is gone, it stops + // heartbeating, and its streams end. Nothing of it settles the claim. + sqlDB, err := aDB.DB() + Expect(err).ToNot(HaveOccurred()) + Expect(sqlDB.Close()).To(Succeed()) + a.stopHeartbeat() + Expect(db.Exec(`UPDATE instances SET last_seen = now() - interval '1 hour' WHERE id = 'replica-a'`).Error).To(Succeed()) + killA() + // The worker loses A and dials B. + Expect(first.Close()).To(Succeed()) + startAgent(b.url) + Eventually(holds(b), "20s", "50ms").Should(BeTrue()) + + // B was there all along. Its loop returns the claim of A to the queue, + // takes it, and drives it on the worker. + startLoop(b, openDB()) + + Eventually(jobStatus(jobID), "40s", "100ms").Should(Equal("completed")) + job, err := store.GetJob(jobID) + Expect(err).ToNot(HaveOccurred()) + Expect(job.Result).To(Equal("recovered")) + Eventually(claims, "10s", "100ms").Should(BeZero()) + Expect(runs.Load()).To(BeEquivalentTo(2), "the job ran again, which is what at-least-once means, and ran to its end once") + }) + }) +}) diff --git a/core/services/carrier/tunnel_integration_test.go b/core/services/carrier/tunnel_integration_test.go index 313dd3417b2d..7d9ec3112a8f 100644 --- a/core/services/carrier/tunnel_integration_test.go +++ b/core/services/carrier/tunnel_integration_test.go @@ -26,8 +26,10 @@ import ( clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging/messagingtest" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/storage" "github.com/mudler/LocalAI/core/services/testutil" "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/LocalAI/core/services/tunnel/slowlink" @@ -93,6 +95,7 @@ type controlPlane struct { installs []workerctl.BackendInstallRequest stops []workerctl.ModelStopRequest renewals int + ensures []workerctl.FileEnsureRequest } func (c *controlPlane) handler() http.Handler { @@ -131,6 +134,14 @@ func (c *controlPlane) handler() http.Handler { c.mu.Unlock() _ = json.NewEncoder(w).Encode(workerctl.OperationReply{Renewed: req.Renew, Completed: req.Complete}) }) + mux.HandleFunc(workerctl.PathOf(workerctl.VerbFilesEnsure), func(w http.ResponseWriter, r *http.Request) { + var req workerctl.FileEnsureRequest + body(r, &req) + c.mu.Lock() + c.ensures = append(c.ensures, req) + c.mu.Unlock() + _ = json.NewEncoder(w).Encode(workerctl.FileEnsureReply{LocalPath: "/cache/" + req.Key}) + }) mux.HandleFunc(workerctl.PathOf(workerctl.VerbBackendList), func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"backends":[]}`) }) @@ -146,6 +157,12 @@ func (c *controlPlane) seenStops() []workerctl.ModelStopRequest { return slices.Clone(c.stops) } +func (c *controlPlane) seenEnsures() []workerctl.FileEnsureRequest { + c.mu.Lock() + defer c.mu.Unlock() + return slices.Clone(c.ensures) +} + func (c *controlPlane) seenRenewals() int { c.mu.Lock() defer c.mu.Unlock() @@ -166,12 +183,21 @@ type replica struct { commands *carrier.Commands files *carrier.Files clients *carrier.Clients + agents *carrier.Agents + queue *carrier.WorkQueue + bus *carrier.Broadcaster + + // selector picks the agent workers, as the set does for the agent control. + selector *nodes.AgentSelector + // stopHeartbeat ends the heartbeat of the row of this replica, as the death + // of its process does. + stopHeartbeat context.CancelFunc } // startReplica starts one frontend: its own tunnel registry, peer link, relay and // carrier set, behind its own HTTP server, over the database that all replicas // share. -func startReplica(ctx context.Context, id string, db *gorm.DB, dsn string, clusterR *cluster.Registry, nodeReg *nodes.NodeRegistry) *replica { +func startReplica(ctx context.Context, id string, db *gorm.DB, dsn string, clusterR *cluster.Registry, nodeReg *nodes.NodeRegistry, tweak ...func(*carrier.TunnelOptions)) *replica { GinkgoHelper() r := &replica{id: id} cred := cluster.NewPeerCredential() @@ -189,21 +215,42 @@ func startReplica(ctx context.Context, id string, db *gorm.DB, dsn string, clust r.url = srv.URL Expect(clusterR.RegisterPeer(ctx, id, "test", 0, "", strings.TrimPrefix(srv.URL, "http://"), cred.Hash())).To(Succeed()) + // The row of the replica stays live while its process lives. + beat, stop := context.WithCancel(ctx) + r.stopHeartbeat = stop + DeferCleanup(stop) + go func() { + t := time.NewTicker(2 * time.Second) + defer t.Stop() + for { + select { + case <-beat.Done(): + return + case <-t.C: + _ = clusterR.Register(beat, id, "test", 0, "") + } + } + }() + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) Expect(err).ToNot(HaveOccurred()) DeferCleanup(fan.Close) - fake := newFakeCarrier(cluster.CarrierTunnel, 1) - set, err := carrier.NewTunnelSet(carrier.TunnelOptions{ + r.selector = nodes.NewAgentSelector(nodeReg, clusterR, id) + options := carrier.TunnelOptions{ Epoch: 1, Dialer: tunnel.NewWorkerDialer(r.tunnels, r.pool), Fanout: fan, - WorkQueue: fake.queue, - Agents: fake.agents, + WorkQueue: jobs.NewClaimQueue(db, fan.Broadcaster), + AgentSelector: r.selector, Registry: nodeReg, InstallTimeout: 30 * time.Second, UpgradeTimeout: 30 * time.Second, Token: registrationToken, - }) + } + for _, t := range tweak { + t(&options) + } + set, err := carrier.NewTunnelSet(options) Expect(err).ToNot(HaveOccurred()) r.set = set r.active = &atomic.Pointer[carrier.Set]{} @@ -211,6 +258,9 @@ func startReplica(ctx context.Context, id string, db *gorm.DB, dsn string, clust r.commands = carrier.NewCommands(r.active) r.files = carrier.NewFiles(r.active) r.clients = carrier.NewClients(r.active) + r.agents = carrier.NewAgents(r.active) + r.queue = carrier.NewWorkQueue(r.active) + r.bus = carrier.NewBroadcaster(r.active) return r } @@ -554,6 +604,49 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { Expect(a.commands.PingNode(nodeID)).To(Succeed()) }) + It("stages a file through shared object storage, with only the verbs on the tunnel, when the set is built for it", func() { + store, err := storage.NewFilesystemStore(GinkgoT().TempDir()) + Expect(err).ToNot(HaveOccurred()) + fm, err := storage.NewFileManager(store, GinkgoT().TempDir()) + Expect(err).ToNot(HaveOccurred()) + s3 := startReplica(ctx, "replica-s3", db, dsn, clusterR, nodeReg, func(o *carrier.TunnelOptions) { + o.S3Staging, o.FileManager = true, fm + }) + attach(s3.url) + Eventually(ownerIs(s3), "15s", "50ms").Should(BeTrue()) + + src := filepath.Join(GinkgoT().TempDir(), "x.bin") + Expect(os.WriteFile(src, []byte("shared bytes"), 0o600)).To(Succeed()) + remote, err := s3.files.EnsureRemote(ctx, nodeID, src, "models/s3/x.bin") + Expect(err).ToNot(HaveOccurred()) + Expect(remote).To(Equal("/cache/models/s3/x.bin")) + + exists, err := store.Exists(ctx, "models/s3/x.bin") + Expect(err).ToNot(HaveOccurred()) + Expect(exists).To(BeTrue(), "the frontend put the file in the shared store") + Expect(control.seenEnsures()).To(Equal([]workerctl.FileEnsureRequest{{Key: "models/s3/x.bin"}})) + Expect(filepath.Join(stagingDir, "models", "s3", "x.bin")).ToNot(BeAnExistingFile(), "no bytes went over the tunnel") + }) + + It("refuses to build a set that lacks a member of this carrier", func() { + whole := carrier.TunnelOptions{ + Dialer: tunnel.NewWorkerDialer(a.tunnels, a.pool), Fanout: &carrier.Fanout{}, + WorkQueue: a.queue, AgentSelector: a.selector, + } + for name, mutate := range map[string]func(*carrier.TunnelOptions){ + "no dialer": func(o *carrier.TunnelOptions) { o.Dialer = nil }, + "no fan-out": func(o *carrier.TunnelOptions) { o.Fanout = nil }, + "no queue": func(o *carrier.TunnelOptions) { o.WorkQueue = nil }, + "no agent selector": func(o *carrier.TunnelOptions) { o.AgentSelector = nil }, + "S3 staging with no manager": func(o *carrier.TunnelOptions) { o.S3Staging = true }, + } { + o := whole + mutate(&o) + _, err := carrier.NewTunnelSet(o) + Expect(err).To(HaveOccurred(), name) + } + }) + It("carries broadcasts between the replicas through the holders", func() { subject := messagingtest.BroadcastRootSubjects["state"] var got atomic.Int32 From 00e222be93301ffb5a91c9cabb82091c107168f9 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:32:00 +0000 Subject: [PATCH 086/162] docs(distributed): describe the claim queue, the agent worker on the tunnel and the S3 stager The contributor guide lists the queue, the agent RPC and the file stager for the tunnel carrier next to those of NATS, and says how the claim queue settles a row, why it has no dead letter, what it does with a cancel and a plain task, and what the shared work queue suite states and leaves out. The list of open items drops the ones that are done. The claim table is migrated at startup on every carrier, so that a change of carrier finds it. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 120 ++++++++++++++++++++++++++------ core/application/distributed.go | 5 ++ 2 files changed, 102 insertions(+), 23 deletions(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index bfc67be80c4f..564045c87ff6 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -17,11 +17,11 @@ files to that spec. | Seam | Interface | Lives in | Today | |---|---|---|---| | Fan-out | `messaging.Broadcaster` | `core/services/messaging` | NATS client | -| Queues | `messaging.WorkQueue` (producer), `messaging.WorkConsumer` (worker), keyed by `messaging.WorkKind` | `core/services/messaging` (`workqueue.go`) | `NewNATSWorkQueue`, `NewNATSWorkConsumer` | -| Control verbs, frontend half | `nodes.NodeCommandSender`, `nodes.FileStager` | `core/services/nodes` | `RemoteUnloaderAdapter`, `S3NATSFileStager` over NATS request/reply; `HTTPFileStager` over HTTP; `TunnelControl` over the HTTP control plane of a worker | +| Queues | `messaging.WorkQueue` (producer), `messaging.WorkConsumer` (worker), keyed by `messaging.WorkKind` | `core/services/messaging` (`workqueue.go`) | `NewNATSWorkQueue`, `NewNATSWorkConsumer`; `jobs.ClaimQueue` and `jobs.ClaimConsumer` (the claim table); `agentworker.Work` (an agent worker on the tunnel) | +| Control verbs, frontend half | `nodes.NodeCommandSender`, `nodes.FileStager` | `core/services/nodes` | `RemoteUnloaderAdapter` over NATS request/reply and `TunnelControl` over the HTTP control plane of a worker; `S3FileStager` over either (`NewS3NATSFileStager`, `NewS3TunnelFileStager`); `HTTPFileStager` over HTTP | | Control verbs, worker half | unexported `controlServer` (`handle`, `handleWithProgress`) | `core/services/worker` | `natsControlServer`, `httpControlServer` | -| Agent RPC, frontend half | `AgentControl` (package `mcp`) | `core/http/endpoints/mcp` | `nodes.NATSAgentControl` | -| Agent RPC, worker half | unexported `agentRPCServer` | `core/cli/agent_worker.go` | `nodes.NATSAgentRPCServer` | +| Agent RPC, frontend half | `AgentControl` (package `mcp`) | `core/http/endpoints/mcp` | `nodes.NATSAgentControl`, `nodes.AgentControlClient` (tunnel) | +| Agent RPC, worker half | unexported `agentRPCServer` on NATS, `agentworker.Handler` on the tunnel | `core/cli/agent_worker.go`, `core/services/agentworker` | `nodes.NATSAgentRPCServer`, `agentworker.Config` and `agentworker.Work` | | Dial to a worker | `nodes.BackendClientFactory`, `nodes.ModelProber`, `nodes.WorkerNetDialerFor` | `core/services/nodes` | direct dial | Request and reply payloads of the control verbs live in @@ -140,12 +140,66 @@ MCP CI with 1 (`startMCPCIConsumer`) and for agent runs with the dispatcher's An empty subject keeps `agent.execute`; an empty queue is kept and makes a plain subscription, as an explicitly empty `LOCALAI_AGENT_QUEUE` always did. +`messagingtest.RunWorkQueueConformance` is the suite every carrier of the +queue passes: one consumer gets a payload and a publisher for its events, +competing consumers split the work once when nothing fails, the kinds stay apart, +`maxInFlight` holds, `Unsubscribe` waits for the running handler, and a payload +above `messaging.MaxWorkPayloadBytes` (1 MiB, the default `max_payload` of NATS) +is refused. It does not state what differs between carriers: the fate of work +whose consumer dies, and how often a failed handler is called again. A handler +must tolerate a repeat. + +### The claim queue + +The tunnel carrier has no broker, so its queue is a table. `jobs.ClaimQueue` +writes a row (`work_claims`) and sends a wake hint on `messaging.SubjectClaimWake`. +`jobs.ClaimConsumer` runs on a frontend replica. It polls every two seconds, +wakes at once on a hint for its kind, and claims into as many slots as it has with +`FOR UPDATE SKIP LOCKED`. The hint is at-most-once, so the poll is what guarantees +that work is found; with the poll alone a unit waits one second on average, and +with the hint the median wait is a few milliseconds. A row stays until the work +has an answer, so `Enqueue` succeeds while nobody consumes. + +A replica claims and not an agent worker, because an agent worker has no database. +`jobs.AgentDriver` is the handler: it picks an agent worker with +`nodes.AgentSelector`, sends the payload to the verb of the kind +(`workerctl.VerbAgentExecute` or `VerbMCPCIRun`) and reads the stream that comes +back. `jobs.DispatchLoop` is a consumer for the three kinds with the driver as its +handler. + +- The outcome of the handler settles the row. `nil` deletes it. An error, or a + panic, releases it with a wait that doubles from two seconds to a minute (the + wait keeps one poison row from being claimed ahead of every newer row on every + tick). An error that wraps `jobs.ErrKeepClaim` leaves it with its replica: the + work ran and its answer could not be recorded, and releasing would run it again. +- There is no dead letter. The only outcomes that release a row are those where + nothing was learned about the work, and failing a job on them would report a + missing connection as the verdict of a worker. +- A claim is settled only by the replica and the attempt that hold it, so a + holder that was reaped and answers late cannot delete the claim of the replica + that took the work over. +- `ReapAbandoned` returns the claims of replicas that are not live. It reads + `cluster.LiveInstanceIDsSQL` and has no age in it: a live replica that holds a + claim for an hour has a slow job. A replica whose own row is not live claims + nothing. +- Delivery is at-least-once. A replica that dies in the middle of a job loses + nothing and the job runs again. NATS loses the job of a consumer that dies. +- `MigratePending` takes the pending rows out of the queue once, for a change to a + carrier that has its own queue. Held rows are driven to their end where they are. +- A plain task (`WorkTask`) has no worker on either carrier. The driver drops the + claim and the job is left to the reaper (`ReapStuckJobs`), as it was on NATS. +- `Dispatcher.Cancel` publishes `jobs..cancel`. On NATS nothing subscribes, as + on master, and a spec pins that. On this carrier the replica that holds the job + ends its stream, the run on the worker ends with its request, and the job is + closed as cancelled. + ## Control verbs Frontend half: `NodeCommandSender` sends the lifecycle verbs; `FileStager` -moves files. `S3NATSFileStager` returns `nodes.ErrNoRoute` when nothing is -listening for the node. `HTTPFileStager` reports connection failures as -ordinary errors. +moves files. `S3FileStager` sends the file verbs over a control link, so the same +stager runs over NATS and over the HTTP control plane of the tunnel, and it +returns `nodes.ErrNoRoute` when no route to the node exists. `HTTPFileStager` +reports connection failures as ordinary errors. `NodeCommandSender` embeds `LoadOperationControl`: the load operation verbs (`InstallBackendOp`, `StopLoadOperation`, `OperationControl` for renewals and @@ -202,6 +256,27 @@ cancellation, and uses the default MCP timeouts when there is none. undecodable body with an `unmarshal error: ` reply. It also serves the node's backend stop as `func(backend string)` and never replies to it. +On the tunnel the frontend picks the worker. `nodes.AgentSelector` lists the agent +nodes that may take work (not pending, not draining; a failed health probe does +not exclude one) and asks which of them hold a tunnel, preferring one that this +replica holds because that call needs no relay. `nodes.AgentControlClient` sends +the request over the control client, with the same contract as NATS: a decoded +reply is returned with a nil error, no agent worker wraps `nodes.ErrNoRoute`, and +only the deadline of ctx applies. A request goes to a second worker (up to three) +only when the failure proves that the first did not start it: no route, a worker +too old to serve the verb, or a worker with no free slot. A broken stream or a +timeout is returned, because the tool may have run. + +`agentworker.Handler` is the worker half on the tunnel: one HTTP server that +binds loopback only and is reached through the tunnel (its tunnel offers the http +tag alone). `agentworker.Work` is its `WorkConsumer`. A run is a streaming +request: the events that the handler publishes are progress lines that name their +subject, the result it publishes becomes the reply line, and a body that ends +without a reply line is a link that broke. A worker that is full answers busy +(`workerctl.WriteBusy`, `workerctl.ErrWorkerBusy`) and the frontend offers the run +to another worker. The frontend publishes a line only if +`nodes.Rebroadcaster` allows the subject for the node type of the worker. + ## Dial `BackendClientFactory.NewClient(nodeID, address, parallel)` builds the gRPC @@ -347,24 +422,23 @@ other seams against a real server, also through Docker. ## Open items for a second carrier -- The agent worker still uses NATS directly for its connection and for agent - events. `agents.NewEventBridge` takes a `Broadcaster`, so the frontend's - bridge goes through the holder. -- `WorkHandler` returns only an error. A carrier whose stream handler must send - a terminal reply derives it from the `jobs..result` event the handler - publishes on `events` (`handleMCPCIJob` does this). -- Not additive: the agent-run consumer (`NATSDispatcher.runDelivery`) ignores - the per-delivery `events` publisher and publishes through the process-wide - `EventBridge`, which is bound to one `Broadcaster`. A second carrier must - change `NATSDispatcher` and `EventBridge`, for example by binding - `handleJob`'s publishes to `events` through a bridge view that shares the - cancel registry. +- `WorkHandler` returns only an error. A carrier whose stream handler must send a + terminal reply derives it from the `jobs..result` event the handler + publishes on `events` (`handleMCPCIJob` does this, and `agentworker.Work` + builds the reply line from it). +- The agent-run consumer (`NATSDispatcher.runDelivery`) binds the events of a run + to the publisher of its delivery through `EventBridge.WithPublisher`, which + shares the cancel registry with the bridge. On NATS that publisher is the NATS + client, so the subjects and payloads are the ones as before. - Agent cancel has no production sender (`EventBridge.CancelExecution` has no - caller), so it is not part of the agent RPC seam. -- The file stager over shared object storage (`S3NATSFileStager`) sends its - requests over NATS. The files verbs exist on the HTTP control plane, but no - frontend stager uses them yet. + caller), so it is not part of the agent RPC seam, and the tunnel has no verb for + it. The cancel of a run on the tunnel is its request. - A worker that registers an address and holds a tunnel (a worker that can use either carrier) is reached by the backend-logs proxy through the host that its address names, and the proxy of the environment is used for it. Only a worker without an address skips the proxy. +- The health monitor reads the presence of each node (`HealthMonitor.UsePresence`) + only while the tunnel is the active carrier: a node that heartbeats and whose + tunnel has been gone longer than `cluster.DefaultReconnectGrace` is demoted with + a status-only change, and is neither promoted nor probed until its tunnel is + back. diff --git a/core/application/distributed.go b/core/application/distributed.go index 3cac082cfac3..aeb0a3386cc4 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -378,6 +378,11 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade if err != nil { return nil, fmt.Errorf("initializing job store: %w", err) } + // The claim queue is the queue of the tunnel carrier. Its table exists on + // every carrier, so that a change of carrier finds it. + if err := jobs.MigrateClaims(context.Background(), authDB); err != nil { + return nil, err + } xlog.Info("Distributed job store initialized") // Initialize job dispatcher From 36ebc33b935fff862bff7bd83be1cf55c1acaaff Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:33:33 +0000 Subject: [PATCH 087/162] docs(messaging): say what the work queue delivers on each carrier WorkQueue and WorkConsumer now state that the claim queue keeps a row until the work has an answer, that NATS delivers at most once and the claim queue at least once, that both refuse a payload above one bound, and that a handler must be idempotent. A spec also checks that the rebroadcaster logs the refusal of a subject outside its list. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/messaging/workqueue.go | 7 ++++- core/services/nodes/rebroadcast_test.go | 34 +++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/core/services/messaging/workqueue.go b/core/services/messaging/workqueue.go index 898fc894b93e..44df2e423e44 100644 --- a/core/services/messaging/workqueue.go +++ b/core/services/messaging/workqueue.go @@ -21,7 +21,12 @@ const ( // apply to the queue. // // The NATS carrier ignores ctx: its publish takes none and returns once the -// message is buffered, so there is nothing for a cancellation to interrupt. +// message is buffered, so there is nothing for a cancellation to interrupt. The +// claim queue (jobs.ClaimQueue) writes a row with ctx, and keeps it until the +// work has an answer, so a nil error from it also means that the work is not +// lost if no consumer exists yet. +// +// Both refuse a payload above MaxWorkPayloadBytes. type WorkQueue interface { Enqueue(ctx context.Context, kind WorkKind, payload any) error } diff --git a/core/services/nodes/rebroadcast_test.go b/core/services/nodes/rebroadcast_test.go index bb2134cb715a..b0a7734ce935 100644 --- a/core/services/nodes/rebroadcast_test.go +++ b/core/services/nodes/rebroadcast_test.go @@ -1,16 +1,38 @@ package nodes import ( + "bytes" "encoding/json" "errors" + "log/slog" + "sync" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/xlog" ) +// syncBuffer is a log destination that handler goroutines can write to. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + // failingBroadcaster refuses every publish. type failingBroadcaster struct{ messaging.Broadcaster } @@ -64,6 +86,18 @@ var _ = Describe("The broadcasts that a worker may ask for", func() { Expect(bus.PublishCount("jobs.j1.cancel")).To(BeZero()) }) + It("logs the refusal of a subject outside the list, so that a worker asking for more is seen", func() { + logs := &syncBuffer{} + xlog.SetLogger(xlog.NewLoggerWithHandler(slog.NewTextHandler(logs, &slog.HandlerOptions{Level: slog.LevelWarn}), xlog.LogLevelWarn)) + DeferCleanup(func() { xlog.SetLogger(xlog.NewLogger(xlog.LogLevel("info"), "text")) }) + + bus := testutil.NewFakeBus() + Expect(NewRebroadcaster(bus).Handle(NodeTypeAgent, "nodes.n1.backend.stop", json.RawMessage(`{}`))).To(BeFalse()) + Expect(logs.String()).To(ContainSubstring("refusing a worker's re-broadcast request")) + Expect(logs.String()).To(ContainSubstring("nodes.n1.backend.stop")) + Expect(bus.PublishCount("nodes.n1.backend.stop")).To(BeZero()) + }) + It("refuses every subject for a backend worker", func() { bus := testutil.NewFakeBus() Expect(NewRebroadcaster(bus).Handle(NodeTypeBackend, "jobs.j1.progress", json.RawMessage(`{}`))).To(BeFalse()) From c9dc8ce3666c8faa8778aeb3abfabc2b5e46e655 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:35:09 +0000 Subject: [PATCH 088/162] style: check the errors that the linter asks about The event bridge calls in the agent dispatcher and the stop of the dispatcher in the agent worker ignore their errors on purpose, as they did before; the blank assignment says so. A field that the S3 stager spec never used is gone. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/agent_worker.go | 2 +- core/services/agents/dispatcher.go | 16 ++++++++-------- .../nodes/file_stager_s3_carriers_test.go | 5 ++--- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 3401cd6cd278..8f62a8c790ef 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -299,7 +299,7 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { xlog.Info("Shutting down agent worker") shutdownCancel() // stop heartbeat loop immediately - dispatcher.Stop() + _ = dispatcher.Stop() mcpTools.CloseAllMCPSessions() regClient.GracefulDeregister(nodeID) return runErr diff --git a/core/services/agents/dispatcher.go b/core/services/agents/dispatcher.go index 528ea7fffe66..b477fc3dd132 100644 --- a/core/services/agents/dispatcher.go +++ b/core/services/agents/dispatcher.go @@ -286,7 +286,7 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent, brid if err != nil { xlog.Error("Failed to load agent config", "agent", evt.AgentName, "error", err) if bridge != nil { - bridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") + _ = bridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") } return } @@ -294,7 +294,7 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent, brid if cfg == nil { xlog.Error("No agent config available", "agent", evt.AgentName) if bridge != nil { - bridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") + _ = bridge.PublishStatus(evt.AgentName, evt.UserID, "error: agent config not found") } return } @@ -333,8 +333,8 @@ func (d *NATSDispatcher) handleJob(ctx context.Context, evt AgentChatEvent, brid if execErr != nil { xlog.Error("Distributed agent execution failed", "agent", evt.AgentName, "error", execErr) if bridge != nil { - bridge.PublishStatus(evt.AgentName, evt.UserID, "error") - bridge.PublishMessage(evt.AgentName, evt.UserID, RoleAgent, + _ = bridge.PublishStatus(evt.AgentName, evt.UserID, "error") + _ = bridge.PublishMessage(evt.AgentName, evt.UserID, RoleAgent, fmt.Sprintf("Agent execution failed: %v", execErr), evt.MessageID+"-error") } return @@ -428,7 +428,7 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent, bridge *EventBri default: return } - bridge.PublishStreamEvent(evt.AgentName, evt.UserID, data) + _ = bridge.PublishStreamEvent(evt.AgentName, evt.UserID, data) }, OnReasoning: func(text string) { // Reasoning is buffered via OnStream @@ -439,7 +439,7 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent, bridge *EventBri OnToolResult: func(name, result string) { // Emit tool_result stream event for real-time UI display if bridge != nil { - bridge.PublishStreamEvent(evt.AgentName, evt.UserID, map[string]any{ + _ = bridge.PublishStreamEvent(evt.AgentName, evt.UserID, map[string]any{ "type": "tool_result", "tool_name": name, "tool_result": result, @@ -462,12 +462,12 @@ func (d *NATSDispatcher) buildNATSCallbacks(evt AgentChatEvent, bridge *EventBri }, OnStatus: func(status string) { if bridge != nil { - bridge.PublishStatus(evt.AgentName, evt.UserID, status) + _ = bridge.PublishStatus(evt.AgentName, evt.UserID, status) } }, OnMessage: func(sender, content, msgID string) { if bridge != nil { - bridge.PublishMessage(evt.AgentName, evt.UserID, sender, content, msgID) + _ = bridge.PublishMessage(evt.AgentName, evt.UserID, sender, content, msgID) } // On agent response, persist the root observable with completion diff --git a/core/services/nodes/file_stager_s3_carriers_test.go b/core/services/nodes/file_stager_s3_carriers_test.go index c6b08d7c4b22..8e08756e458a 100644 --- a/core/services/nodes/file_stager_s3_carriers_test.go +++ b/core/services/nodes/file_stager_s3_carriers_test.go @@ -78,9 +78,8 @@ func (r *natsRig) Request(subject string, data []byte, _ time.Duration) ([]byte, var _ = Describe("The S3 file stager over each carrier of the file verbs", func() { type rig struct { - name string - build func(fm *storage.FileManager, w *scriptedWorker) (FileStager, func(noRoute bool)) - release func(FileStager) RequestFileReleaser + name string + build func(fm *storage.FileManager, w *scriptedWorker) (FileStager, func(noRoute bool)) } rigs := []rig{ From 8a602a568265fed37a1cd5b1619b3edb084a710e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:51:45 +0000 Subject: [PATCH 089/162] fix(tunnel): bound the relay frame by the budget of the caller The dialer wrote the relay request and read its reply before any deadline was armed. An owner replica whose session still answered the keepalive but whose relay handler was stuck held the caller, a goroutine and a stream for as long as the link lived. A dial with a 500 ms budget never returned, and the lease renewal that depends on it stopped. The relay hop now arms the deadline of the caller before the frame, and a watcher closes the stream when the context ends, which covers a caller that cancels without a deadline. The handshake with the worker gets the same watcher. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/dialer.go | 32 +++++++++++++++ core/services/tunnel/dialer_test.go | 61 +++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/core/services/tunnel/dialer.go b/core/services/tunnel/dialer.go index 2df40ff23acd..a60846dbd00a 100644 --- a/core/services/tunnel/dialer.go +++ b/core/services/tunnel/dialer.go @@ -1,6 +1,7 @@ package tunnel import ( + "cmp" "context" "errors" "fmt" @@ -147,9 +148,24 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) } + // The relay frame and its reply run under the budget of the caller. An owner + // whose session still answers the keepalive and whose handler is stuck would + // otherwise hold the caller, a goroutine and a stream for as long as the link + // lives. The deadline covers a caller that stated one, and the watcher covers + // a caller that cancels without one. + if err := stream.SetDeadline(handshakeDeadline(ctx)); err != nil { + _ = stream.Close() + return nil, routeFailure(nodeID, fmt.Errorf("arming the relay deadline: %w", err)) + } + release := closeWhenDone(ctx, stream) + defer func() { _ = release() }() + // The remaining time of the caller, so that the owner can bound its own open. if err := WriteRelayRequest(stream, nodeID, relayLane(o), remainingBudget(ctx)); err != nil { _ = stream.Close() + if blamed := callerRanOut(ctx); blamed != nil { + return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, blamed)) + } return nil, routeFailure(nodeID, fmt.Errorf("naming the node on a stream to replica %q: %w", owner, err)) } if err := ReadRelayReply(stream); err != nil { @@ -187,6 +203,8 @@ func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, t _ = stream.Close() return nil, routeFailure(nodeID, fmt.Errorf("arming the handshake deadline: %w", err)) } + release := closeWhenDone(ctx, stream) + defer func() { _ = release() }() if err := WriteStreamRequest(stream, tag, target); err != nil { // The tunnel broke under the request. Nothing was asked of the worker and @@ -217,10 +235,24 @@ func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, t _ = stream.Close() return nil, routeFailure(nodeID, fmt.Errorf("clearing the handshake deadline: %w", err)) } + if release() { + // The caller gave up in the last moment and the watch closed the stream. + return nil, routeFailure(nodeID, fmt.Errorf("opening %q: the budget of the caller ran out: %w", tag, cmp.Or(callerRanOut(ctx), context.Canceled))) + } xlog.Debug("opened a tunnelled stream to a worker", "node", nodeID, "tag", tag, "target", target) return stream, nil } +// closeWhenDone closes the stream when ctx ends, so that a read or a write that +// is parked on a peer which says nothing returns at once. The deadline of the +// socket covers a caller that stated one; this covers a caller that only +// cancels. The returned function stops the watch and must be called when the +// exchange is over. It reports whether the watch had already closed the stream. +func closeWhenDone(ctx context.Context, stream net.Conn) (release func() (fired bool)) { + stop := context.AfterFunc(ctx, func() { _ = stream.Close() }) + return func() bool { return !stop() } +} + // handshakeDeadline is when the handshake must be done: the deadline of the caller // when it has one and it is sooner, and the backstop otherwise. func handshakeDeadline(ctx context.Context) time.Time { diff --git a/core/services/tunnel/dialer_test.go b/core/services/tunnel/dialer_test.go index 751ad7a708a1..d5a007623a91 100644 --- a/core/services/tunnel/dialer_test.go +++ b/core/services/tunnel/dialer_test.go @@ -263,6 +263,42 @@ var _ = Describe("The worker dialer", func() { Expect(worker.seen()).To(BeEmpty(), "nothing may have crossed the inference lane") }) + It("gives up on an owner that took the relay frame and never answers, when the budget of the caller runs out", func() { + wedged := &wedgedPeers{} + dialer := tunnel.NewWorkerDialer(tunnels, wedged) + short, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + defer cancel() + + done := make(chan error, 1) + go func() { + _, err := dialer.Dial(short, "w1", tunnel.StreamTagGRPC, "x:1") + done <- err + }() + var err error + Eventually(done, 2*time.Second).Should(Receive(&err), "a wedged owner must not hold the caller past its budget") + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "an expired budget is not a missing route") + Eventually(wedged.farClosed, 2*time.Second).Should(BeTrue(), "the stream to the owner must be closed") + }) + + It("closes the stream to the owner when the caller cancels and stated no deadline", func() { + wedged := &wedgedPeers{} + dialer := tunnel.NewWorkerDialer(tunnels, wedged) + cancelable, cancel := context.WithCancel(ctx) + + done := make(chan error, 1) + go func() { + _, err := dialer.Dial(cancelable, "w1", tunnel.StreamTagGRPC, "x:1") + done <- err + }() + time.Sleep(200 * time.Millisecond) + cancel() + var err error + Eventually(done, 2*time.Second).Should(Receive(&err)) + Expect(errors.Is(err, context.Canceled)).To(BeTrue(), "%v", err) + Eventually(wedged.farClosed, 2*time.Second).Should(BeTrue()) + }) + It("states the remaining budget of the caller in the relay request", func() { dialer := tunnel.NewWorkerDialer(tunnels, peers) budgeted, cancel := context.WithTimeout(ctx, 20*time.Second) @@ -439,3 +475,28 @@ func tunnelUnreachable() error { Expect(err).To(MatchError(tunnel.ErrPeerUnreachable)) return err } + +// wedgedPeers hands out a stream to an owner that reads the relay frame and +// then says nothing, the way a replica does when its session still answers the +// keepalive and its handler is stuck. +type wedgedPeers struct { + mu sync.Mutex + end bool +} + +func (w *wedgedPeers) Open(context.Context, string) (net.Conn, error) { + near, far := net.Pipe() + go func() { + _, _ = io.Copy(io.Discard, far) + w.mu.Lock() + w.end = true + w.mu.Unlock() + }() + return near, nil +} + +func (w *wedgedPeers) farClosed() bool { + w.mu.Lock() + defer w.mu.Unlock() + return w.end +} From b5a4cad400ff4b61d5a6c5a577dcbbde3dc8a375 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:54:01 +0000 Subject: [PATCH 090/162] fix(tunnel): let a caller leave the queue for a peer dial when its budget ends PeerPool.Open held a mutex of the peer through the dial and ignored the context while it waited. A caller with a one second budget that arrived behind a caller with none returned after almost ten seconds. The per-peer serialisation is now a gate that a caller takes with a select on its context. The lock that guards the cached session is separate and is never held across a dial, so Close does not wait behind one either. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/peerlink.go | 78 ++++++++++++++++++++------ core/services/tunnel/peerpool_test.go | 80 +++++++++++++++++++++++++++ 2 files changed, 142 insertions(+), 16 deletions(-) create mode 100644 core/services/tunnel/peerpool_test.go diff --git a/core/services/tunnel/peerlink.go b/core/services/tunnel/peerlink.go index 3eded8566ad7..4756f6b512b2 100644 --- a/core/services/tunnel/peerlink.go +++ b/core/services/tunnel/peerlink.go @@ -84,9 +84,7 @@ func rejectedPeer(peerID string, cause error) error { // this process and says nothing about whether the peer exists. var ErrPoolClosed = errors.New("tunnel: peer pool is closed") -// peerLinkHandshakeTimeout bounds the websocket upgrade of a peer link. It also -// bounds how long Close waits behind a dial in progress, which holds the lock of -// its peer. +// peerLinkHandshakeTimeout bounds the websocket upgrade of a peer link. const peerLinkHandshakeTimeout = 10 * time.Second // PeerPool dials peer replicas and keeps one multiplexed session for each. @@ -120,14 +118,47 @@ type PeerPool struct { // the alias keeps the callers of this package from importing both. type PeerCredential = cluster.PeerCredential -// peerLink is the cached session for one peer and the lock that serialises -// dialling it. The lock is per peer, so a dial to one peer that hangs does not -// hold up the opens to another. +// peerLink is the cached session for one peer and the gate that serialises +// dialling it. The gate is per peer, so a dial to one peer that hangs does not +// hold up the opens to another. It is a channel and not a mutex, so that a caller +// waiting behind a dial leaves when its own context ends: a lock cannot be +// abandoned, and a caller with a short budget would wait out the dial of a caller +// that has none. type peerLink struct { + gate chan struct{} + + // mu guards sess alone and is never held across a dial or an open, so that + // Close does not wait behind either. mu sync.Mutex sess *Session } +func newPeerLink() *peerLink { return &peerLink{gate: make(chan struct{}, 1)} } + +// enter takes the gate, or reports why the caller left without it. +func (l *peerLink) enter(ctx context.Context) error { + select { + case l.gate <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (l *peerLink) leave() { <-l.gate } + +func (l *peerLink) session() *Session { + l.mu.Lock() + defer l.mu.Unlock() + return l.sess +} + +func (l *peerLink) setSession(s *Session) { + l.mu.Lock() + defer l.mu.Unlock() + l.sess = s +} + // NewPeerPool returns a pool that dials peers as selfID. cred must be the value // that the membership loop published the hash of, which makes a pool that // presents one secret and a row that holds the hash of another unstateable. @@ -159,11 +190,13 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { return nil, err } - l.mu.Lock() - defer l.mu.Unlock() + if err := l.enter(ctx); err != nil { + return nil, err + } + defer l.leave() - if l.sess != nil { - st, err := l.sess.OpenStream(ctx) + if sess := l.session(); sess != nil { + st, err := sess.OpenStream(ctx) if err == nil { return st, nil } @@ -174,8 +207,8 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { } // A session that died between calls is the common case. xlog.Debug("tunnel peer link session unusable, dialling again", "peer", peerID, "error", err) - _ = l.sess.Close() - l.sess = nil + _ = sess.Close() + l.setSession(nil) } sess, err := p.dial(ctx, peerID) @@ -200,7 +233,14 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { } return nil, unreachablePeer(peerID, err) } - l.sess = sess + l.setSession(sess) + // Close ran while this open held the gate and did not see the new session. + if p.isClosed() { + _ = sess.Close() + l.setSession(nil) + _ = st.Close() + return nil, ErrPoolClosed + } return st, nil } @@ -215,7 +255,7 @@ func (p *PeerPool) link(peerID string) (*peerLink, error) { } l, ok := p.links[peerID] if !ok { - l = &peerLink{} + l = newPeerLink() p.links[peerID] = l } return l, nil @@ -294,6 +334,12 @@ func (p *PeerPool) dial(ctx context.Context, peerID string) (*Session, error) { return sess, nil } +func (p *PeerPool) isClosed() bool { + p.mu.Lock() + defer p.mu.Unlock() + return p.closed +} + // Close closes every cached session. It can be called twice, and an Open after it // reports ErrPoolClosed. func (p *PeerPool) Close() { @@ -307,8 +353,8 @@ func (p *PeerPool) Close() { p.links = nil p.mu.Unlock() - // Each session closes under the lock of its own peer, so that closing the pool - // cannot deadlock against an Open that is dialling and about to take p.mu. + // Each session closes under the small lock of its own peer. It is never held + // across a dial, so closing the pool does not wait behind one. for _, l := range links { l.mu.Lock() if l.sess != nil { diff --git a/core/services/tunnel/peerpool_test.go b/core/services/tunnel/peerpool_test.go new file mode 100644 index 000000000000..e85f5bd1b7ad --- /dev/null +++ b/core/services/tunnel/peerpool_test.go @@ -0,0 +1,80 @@ +package tunnel_test + +import ( + "context" + "errors" + "net" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("The peer pool", func() { + // silentListener accepts connections and never answers them, which is what a + // peer looks like while its websocket upgrade hangs. + silentListener := func() string { + GinkgoHelper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + var ( + mu sync.Mutex + conns []net.Conn + ) + go func() { + for { + c, err := lis.Accept() + if err != nil { + return + } + mu.Lock() + conns = append(conns, c) + mu.Unlock() + } + }() + DeferCleanup(func() { + _ = lis.Close() + mu.Lock() + defer mu.Unlock() + for _, c := range conns { + _ = c.Close() + } + }) + return lis.Addr().String() + } + + It("does not hold a caller with a short budget behind a caller that has none", func() { + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg := cluster.NewRegistry(db) + Expect(reg.RegisterPeer(ctx, "slow", "test", 0, "", silentListener(), cluster.NewPeerCredential().Hash())).To(Succeed()) + pool := tunnel.NewPeerPool("replica-a", cluster.NewPeerCredential(), reg) + DeferCleanup(pool.Close) + + // The first caller states no deadline and sits in the dial. + first := make(chan error, 1) + go func() { + _, err := pool.Open(ctx, "slow") + first <- err + }() + time.Sleep(300 * time.Millisecond) + + short, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + defer cancel() + start := time.Now() + _, err := pool.Open(short, "slow") + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "%v", err) + Expect(time.Since(start)).To(BeNumerically("<", 2*time.Second), "the caller waited on a lock and not on its own budget") + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeFalse(), "an expired budget says nothing about the peer") + + pool.Close() + Eventually(first, 15*time.Second).Should(Receive()) + }) +}) From 915be99f97b0ef6b1a9754bb43270d6b278a7633 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 02:56:36 +0000 Subject: [PATCH 091/162] fix(grpc): report a dial that has not finished as a failure of the transport gRPC gives the deadline of a call to the call and not to the dialer. When a health probe ran out of time while the dial through a tunnel was still running, the client had recorded no dial error. The router read that as a dead backend and removed the row of a model that was loaded, and the health monitor counted a miss. The client now notes when a dial starts, and LastDialError reports ErrDialPending until the dial ends. A dial that succeeds clears it. Specs cover the client with a dialer that is stuck, the router with a probe that gives up during the dial, and the health monitor. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/transport_failure_test.go | 75 +++++++++++++++++++ pkg/grpc/backend.go | 6 ++ pkg/grpc/client.go | 24 +++++- pkg/grpc/dialer_test.go | 30 ++++++++ 4 files changed, 133 insertions(+), 2 deletions(-) diff --git a/core/services/nodes/transport_failure_test.go b/core/services/nodes/transport_failure_test.go index 85146703725b..1ff0b913aef9 100644 --- a/core/services/nodes/transport_failure_test.go +++ b/core/services/nodes/transport_failure_test.go @@ -4,7 +4,10 @@ import ( "context" "errors" "fmt" + "io" + "net" "runtime" + "sync" "time" . "github.com/onsi/ginkgo/v2" @@ -141,6 +144,15 @@ var _ = Describe("The health monitor and a probe that learned nothing", func() { Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) }) + It("does not count a probe that ran out of time while the dial was still running", func() { + store, factory, hm := build() + factory.setClient("127.0.0.1:50053", &fakeBackendClient{err: context.DeadlineExceeded, dialErr: grpc.ErrDialPending}) + for range perModelMissThreshold * 3 { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) + }) + It("still reaps after the threshold when the host answered that the backend is gone", func() { store, factory, hm := build() factory.setClient("127.0.0.1:50053", &fakeBackendClient{err: fmt.Errorf("unavailable"), dialErr: backendAnswer{errors.New("gone")}}) @@ -201,3 +213,66 @@ var _ = Describe("The router and a replica it cannot reach", func() { Expect(reg.removeCalls).To(ContainElement("n-old:m")) }) }) + +// firstThenStub hands the first client it is asked for to the probe of a cached +// replica and the stub to every client after it. +type firstThenStub struct { + mu sync.Mutex + first grpc.Backend + stub *stubBackend +} + +func (f *firstThenStub) NewClient(_, _ string, _ bool) grpc.Backend { + f.mu.Lock() + defer f.mu.Unlock() + if f.first != nil { + c := f.first + f.first = nil + return c + } + return f.stub +} + +// slowDialProbe is a stub whose health check goes through a real gRPC client +// with a dialer that is still running when the probe gives up. +type slowDialProbe struct { + *stubBackend + real grpc.Backend +} + +func (p *slowDialProbe) HealthCheck(ctx context.Context) (bool, error) { + return p.real.HealthCheck(ctx) +} +func (p *slowDialProbe) LastDialError() error { return grpc.LastDialErrorOf(p.real) } +func (p *slowDialProbe) Close() error { + if c, ok := p.real.(io.Closer); ok { + return c.Close() + } + return nil +} + +var _ = Describe("The router and a transport that is slow", func() { + It("keeps the row of a replica whose probe ran out of time while the dial was still running", func() { + release := make(chan struct{}) + DeferCleanup(func() { close(release) }) + real := grpc.NewClientWithDialer("10.0.0.70:9001", false, nil, false, "", func(context.Context, string) (net.Conn, error) { + <-release + return nil, errors.New("released") + }) + stub := &stubBackend{healthResult: true, loadResult: &pb.Result{Success: true}} + reg := &fakeModelRouter{ + findAndLockNode: &BackendNode{ID: "n-old", Name: "old-node", Address: "10.0.0.70:50051"}, + findAndLockNM: &NodeModel{NodeID: "n-old", ModelName: "m", Address: "10.0.0.70:9001"}, + findIdleNode: &BackendNode{ID: "n-new", Name: "new-node", Address: "10.0.0.71:50051"}, + } + unloader := &fakeUnloader{installReply: &workerctl.BackendInstallReply{Success: true, Address: "10.0.0.71:9001"}} + router := NewSmartRouter(reg, SmartRouterOptions{Unloader: unloader, ClientFactory: &firstThenStub{ + first: &slowDialProbe{stubBackend: stub, real: real}, + stub: stub, + }}) + + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) + Expect(err).ToNot(HaveOccurred()) + Expect(reg.removeCalls).ToNot(ContainElement("n-old:m"), "a transport that is slow is not a replica that is gone") + }) +}) diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 9f52f0ea3087..3d12c1106ac6 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -46,6 +46,7 @@ func NewClientWithDialer(address string, parallel bool, wd WatchDog, enableWatch // directly, which is the bypass this constructor exists to close. c := buildClient(address, parallel, wd, enableWatchDog, token) c.dialer = func(ctx context.Context, addr string) (net.Conn, error) { + c.startDial() conn, err := dialer(ctx, addr) c.recordDialErr(err) return conn, err @@ -62,6 +63,11 @@ type DialErrorReporter interface { LastDialError() error } +// ErrDialPending is what LastDialError reports while a dial has started and has +// not finished. It is a failure of the transport and not an answer of the host: +// the caller knows nothing about the backend yet. +var ErrDialPending = errors.New("grpc: the dial to the backend has not finished") + // BackendAnswer is implemented by the error of a custom dialer when the dial // reached the host of the backend and the host answered. A refusal of the worker // to open a stream to a backend that is gone is one: it is evidence about the diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index e48f60cd43cd..36567743ff4b 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -44,8 +44,9 @@ type Client struct { // guards inFlight and is taken on every call, and a dialer runs under the // machinery of gRPC, where taking it again is not something this type can // reason about. - dialErrMu sync.Mutex - lastDialErr error + dialErrMu sync.Mutex + dialsPending int + lastDialErr error sync.Mutex opMutex sync.Mutex @@ -1469,17 +1470,36 @@ func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts .. // its row survives one more round. A dial that succeeded clears the value, and // a failure of the transport reads as a failure of the backend, which is what // happened before the dialer existed. Neither is a new hazard. +// +// A dial that has started and not finished reads as ErrDialPending when no failure +// is recorded. gRPC gives the deadline of a call to the call and not to the +// dialer, so a call can time out while the dial still runs. Without this state +// that caller finds no error of the dial and reads a slow transport as a backend +// that is gone. func (c *Client) LastDialError() error { c.dialErrMu.Lock() defer c.dialErrMu.Unlock() + if c.lastDialErr == nil && c.dialsPending > 0 { + return ErrDialPending + } return c.lastDialErr } +// startDial notes that a dial began. The matching recordDialErr ends it. +func (c *Client) startDial() { + c.dialErrMu.Lock() + c.dialsPending++ + c.dialErrMu.Unlock() +} + // recordDialErr stores the outcome of one dial. A success clears the failure // before it, so that a client that recovered does not keep reporting an error // that describes nothing. func (c *Client) recordDialErr(err error) { c.dialErrMu.Lock() + if c.dialsPending > 0 { + c.dialsPending-- + } c.lastDialErr = err c.dialErrMu.Unlock() } diff --git a/pkg/grpc/dialer_test.go b/pkg/grpc/dialer_test.go index 74e88be55a31..a67606434061 100644 --- a/pkg/grpc/dialer_test.go +++ b/pkg/grpc/dialer_test.go @@ -4,7 +4,9 @@ import ( "context" "errors" "net" + "sync" "sync/atomic" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -60,6 +62,34 @@ var _ = Describe("A client with its own dialer", func() { Expect(errors.Is(TransportFailureOf(client), boom)).To(BeTrue()) }) + It("reads a dial that has not finished as a failure of the transport, and clears it when the dial succeeds", func() { + // gRPC hands the deadline of a call to the call and not to the dialer, so a + // call can time out while the dial is still running. The caller then has no + // error of the dial to read, and a slow transport looks like a dead backend. + release := make(chan struct{}) + var once sync.Once + DeferCleanup(func() { once.Do(func() { close(release) }) }) + // The dialer ignores its context, as a dial that is stuck in the middle of a + // handshake does until the handshake ends. + client := NewClientWithDialer("x:1", false, nil, false, "", func(_ context.Context, _ string) (net.Conn, error) { + <-release + var d net.Dialer + return d.DialContext(context.Background(), "tcp", addr) + }) + + short, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + ok, _ := client.HealthCheck(short) + Expect(ok).To(BeFalse()) + Expect(TransportFailureOf(client)).ToNot(BeNil(), "the dial is still running; nothing is known about the backend") + + once.Do(func() { close(release) }) + Eventually(func() error { return TransportFailureOf(client) }, 5*time.Second).Should(BeNil()) + ok, err := client.HealthCheck(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + }) + It("clears the error when the next dial succeeds", func() { var fail atomic.Bool fail.Store(true) From 4be4e4500aa3cebec3f6625ef3fc186e270d0b05 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 03:05:47 +0000 Subject: [PATCH 092/162] fix(tunnel): keep a slow or broken path apart from a missing route routeFailure turned every error it did not list into ErrNoRoute. A worker that took a stream and said nothing, an owner link that died before it answered the relay frame, a worker that answered with garbage and a database that could not name the owner all became a missing route, and a scheduler demotes a worker on that. A socket or multiplexer timeout is now ErrTransport. A reply that breaks the protocol is ErrProtocol. A failed owner lookup is ErrInfrastructure. A relay hop that breaks or goes quiet before the owner answers is ErrPeerUnreachable. ErrNoRoute stays for the routing facts: no replica holds the tunnel, this replica cannot relay, the owner could not open a stream, the worker could not serve one, or the session ended under the request. An error that nobody classified is returned as it is and demotes no one. The handshake backstop is a variable so that a spec can reach it. The shared control conformance suite gets a case for a path that fails and one for an infrastructure that fails. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 9 +- .../nodes/control_tunnel_conformance_test.go | 18 ++++ ...load_operation_control_conformance_test.go | 31 ++++++ core/services/tunnel/dialer.go | 77 ++++++++++++--- core/services/tunnel/dialer_test.go | 95 ++++++++++++++++++- core/services/tunnel/export_test.go | 11 +++ core/services/tunnel/proto.go | 6 +- core/services/tunnel/relay.go | 4 +- core/services/tunnel/route.go | 76 +++++++++++++-- core/services/tunnel/route_test.go | 60 ++++++++++++ 10 files changed, 362 insertions(+), 25 deletions(-) create mode 100644 core/services/tunnel/route_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 564045c87ff6..770a310f97ce 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -329,7 +329,14 @@ identity and is no `ErrNoRoute`. A peer that does not answer is `tunnel.ErrPeerUnreachable` and is not `ErrNoRoute`: a link between two frontends says nothing about a worker, and the scheduler demotes a worker on `ErrNoRoute`. A bulk lane that is down is `tunnel.ErrNoBulkSession` and is not `ErrNoRoute` -either. The budget of the caller is checked first and is never a route failure. +either. A path that is slow (a socket deadline, a worker that takes the stream and +says nothing) is `tunnel.ErrTransport`. A reply that breaks the protocol is +`tunnel.ErrProtocol`. A database that cannot say who holds the tunnel is +`tunnel.ErrInfrastructure`. An owner link that breaks before it answers is +`tunnel.ErrPeerUnreachable`. None of them is `ErrNoRoute`, and a failure that +`routeFailure` does not recognise is none of them either: only the listed routing +facts demote a worker. The budget of the caller is checked first and is never a +route failure. A dial can ask for the bulk lane (`tunnel.WithBulkLane`) and then never falls back to the lane of model calls. The file stager asks for it. diff --git a/core/services/nodes/control_tunnel_conformance_test.go b/core/services/nodes/control_tunnel_conformance_test.go index 2f4e1c00914f..13048695caea 100644 --- a/core/services/nodes/control_tunnel_conformance_test.go +++ b/core/services/nodes/control_tunnel_conformance_test.go @@ -4,6 +4,7 @@ package nodes import ( "context" "encoding/json" + "errors" "fmt" "io" "net" @@ -75,6 +76,23 @@ func (h *httpLoadOperationHarness) NoRoute() { func (h *httpLoadOperationHarness) TimesOut() { h.setMode("hang") } +// The carrier maps the errors of the tunnel before the control client sees them, +// and the tunnel keeps these two apart from a missing route. What reaches this +// package is a plain error that carries no ErrNoRoute. +func (h *httpLoadOperationHarness) PathFails() { + h.setDialErr(errors.New("tunnel: peer \"replica-b\" unreachable: the link ended")) +} + +func (h *httpLoadOperationHarness) InfrastructureFails() { + h.setDialErr(errors.New("tunnel: the cluster database could not answer")) +} + +func (h *httpLoadOperationHarness) setDialErr(err error) { + h.mu.Lock() + defer h.mu.Unlock() + h.dialErr = err +} + func (h *httpLoadOperationHarness) WorkerRefuses() { h.setMode("refuse") } func (h *httpLoadOperationHarness) WorkerAnswers() { h.setMode("answer") } diff --git a/core/services/nodes/load_operation_control_conformance_test.go b/core/services/nodes/load_operation_control_conformance_test.go index de5bc75b2460..59c08a53e3ac 100644 --- a/core/services/nodes/load_operation_control_conformance_test.go +++ b/core/services/nodes/load_operation_control_conformance_test.go @@ -36,6 +36,14 @@ type loadOperationHarness interface { NoRoute() // TimesOut makes every later call get no reply in time. TimesOut() + // PathFails makes every later call fail on the way to the worker with a + // failure that is not a missing route: a link between two frontends that + // dies, a connection of the bus that is closed. + PathFails() + // InfrastructureFails makes every later call fail because something the + // carrier depends on, and that is not the worker, does not answer: the + // database that names the owner of a tunnel, the server of the bus. + InfrastructureFails() // WorkerRefuses makes the worker answer every call with its own refusal. WorkerRefuses() // WorkerAnswers makes the worker answer every call with success. @@ -84,6 +92,18 @@ func (h *natsLoadOperationHarness) TimesOut() { } } +func (h *natsLoadOperationHarness) PathFails() { + for _, s := range h.subjects() { + h.mc.scriptErr(s, nats.ErrConnectionClosed) + } +} + +func (h *natsLoadOperationHarness) InfrastructureFails() { + for _, s := range h.subjects() { + h.mc.scriptErr(s, errors.New("the server of the bus does not answer")) + } +} + func (h *natsLoadOperationHarness) WorkerRefuses() { const node = conformanceNode h.mc.scriptReply(messaging.SubjectNodeBackendInstall(node), workerctl.BackendInstallReply{Success: false, Error: "disk full"}) @@ -187,6 +207,17 @@ var _ = Describe("LoadOperationControl conformance", func() { } }) + It("does not report a broken path or an infrastructure failure as ErrNoRoute", func() { + for name, fail := range map[string]func(){"a path that fails": h.PathFails, "an infrastructure that fails": h.InfrastructureFails} { + fail() + for _, c := range calls() { + err := c.run() + Expect(err).To(HaveOccurred(), "%s: %s", name, c.verb) + Expect(errors.Is(err, ErrNoRoute)).To(BeFalse(), "%s: %s: %v", name, c.verb, err) + } + } + }) + It("does not report the budget of the caller running out as ErrNoRoute", func() { h.TimesOut() ctx, cancel := context.WithCancel(context.Background()) diff --git a/core/services/tunnel/dialer.go b/core/services/tunnel/dialer.go index a60846dbd00a..8a9462ec755a 100644 --- a/core/services/tunnel/dialer.go +++ b/core/services/tunnel/dialer.go @@ -23,7 +23,7 @@ type PeerOpener interface { // worker or a relay that accepted a stream and then said nothing would otherwise // park the caller until the keepalive of the session killed it. A deadline of the // caller is used instead when it is the shorter of the two. -const dialHandshakeTimeout = 15 * time.Second +var dialHandshakeTimeout = 15 * time.Second // WorkerDialer opens connections to a worker through its tunnel, wherever in the // deployment that tunnel is held. @@ -120,7 +120,29 @@ func (d *WorkerDialer) GRPCDialerFor(nodeID string) func(ctx context.Context, ad } // relay opens the stream through the replica that holds the tunnel of the worker. +// +// An owner that refuses with ErrNotOwner says that the worker left it after the +// lookup, which is what a worker that reconnects to another replica does. The +// lookup is made once more, within the budget of the caller, and the second +// answer is final: a third attempt would turn a stale row into a loop. func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o dialOptions) (net.Conn, error) { + stream, err := d.relayOnce(ctx, nodeID, o, false) + if err != nil { + return nil, routeFailure(nodeID, err) + } + return d.handshake(ctx, stream, nodeID, tag, target) +} + +// errOwnerRefused marks an ErrNotOwner that an owning replica wrote, as opposed +// to one that this replica reached by itself. Only the first is worth a second +// lookup. +var errOwnerRefused = errors.New("tunnel: the owner refused the stream") + +// relayOnce looks the owner up and opens the stream through it, up to the end of +// the relay reply. It returns the cause of a failure and leaves its +// classification to the caller. again says that this is the second lookup, in +// which the owner may be this replica: the worker moved here in the meantime. +func (d *WorkerDialer) relayOnce(ctx context.Context, nodeID string, o dialOptions, again bool) (net.Conn, error) { // Owner and not the row: a row outlives its owner by up to a liveness window, // and a dial to what the unjoined read returns would reach a process that is // gone and report the worker as unreachable instead of absent. @@ -128,16 +150,29 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o if err != nil { // ErrNoConnection is the ordinary answer for a worker that has not dialled // in yet, and it must not get out: routeFailure keeps it in the message. - return nil, routeFailure(nodeID, err) + // A database that did not answer is another condition, and it says nothing + // about the worker. + if !isAbsenceClaim(err) && callerRanOut(ctx) == nil { + err = fmt.Errorf("%w: looking up the owner: %w", ErrInfrastructure, err) + } + return nil, err } if owner == d.tunnels.selfID { + if again { + // The worker moved to this replica while the first owner was refusing. + var openOpts []OpenOption + if o.noFallback { + openOpts = append(openOpts, WithoutFallback()) + } + return d.tunnels.Open(ctx, nodeID, o.lane, openOpts...) + } // The table names this replica, and the registry says that the tunnel is // not held here, so the attachment went away between the claim and now. // Relaying would send the request into this same process. - return nil, routeFailure(nodeID, fmt.Errorf("the connection row names this replica, which no longer holds the tunnel: %w", ErrNotOwner)) + return nil, fmt.Errorf("the connection row names this replica, which no longer holds the tunnel: %w", ErrNotOwner) } if d.peers == nil { - return nil, routeFailure(nodeID, fmt.Errorf("the tunnel is held by replica %q: %w", owner, ErrNoRelayPath)) + return nil, fmt.Errorf("the tunnel is held by replica %q: %w", owner, ErrNoRelayPath) } stream, err := d.peers.Open(ctx, owner) @@ -145,7 +180,7 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o // ErrPeerUnreachable and ErrPoolClosed keep their identity. An owner swept // between the lookup and this dial is an absence claim about a replica, and // routeFailure withholds it. - return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) + return nil, fmt.Errorf("through replica %q: %w", owner, err) } // The relay frame and its reply run under the budget of the caller. An owner @@ -155,7 +190,7 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o // a caller that cancels without one. if err := stream.SetDeadline(handshakeDeadline(ctx)); err != nil { _ = stream.Close() - return nil, routeFailure(nodeID, fmt.Errorf("arming the relay deadline: %w", err)) + return nil, fmt.Errorf("arming the relay deadline: %w", err) } release := closeWhenDone(ctx, stream) defer func() { _ = release() }() @@ -164,20 +199,40 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o if err := WriteRelayRequest(stream, nodeID, relayLane(o), remainingBudget(ctx)); err != nil { _ = stream.Close() if blamed := callerRanOut(ctx); blamed != nil { - return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, blamed)) + return nil, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, blamed) } - return nil, routeFailure(nodeID, fmt.Errorf("naming the node on a stream to replica %q: %w", owner, err)) + return nil, relayHopFailure(owner, "naming the node on a stream", err) } if err := ReadRelayReply(stream); err != nil { // A refusal of the owning replica and not of the worker: the owner would // not relay, which is a route that does not exist. _ = stream.Close() if blamed := callerRanOut(ctx); blamed != nil { - return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, blamed)) + return nil, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, blamed) + } + err = relayHopFailure(owner, "reading the relay reply", err) + if errors.Is(err, ErrNotOwner) { + err = fmt.Errorf("%w: %w", errOwnerRefused, err) } - return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) + return nil, err } - return d.handshake(ctx, stream, nodeID, tag, target) + if release() { + return nil, fmt.Errorf("through replica %q: the budget of the caller ran out: %w", owner, cmp.Or(callerRanOut(ctx), context.Canceled)) + } + return stream, nil +} + +// relayHopFailure classifies a failure on the stream to the owning replica. A +// refusal that the owner wrote, and a reply that breaks the protocol, keep their +// identity. Everything else is the link breaking or going quiet before the owner +// answered, which is a fact about two frontends and not about the worker. +func relayHopFailure(owner, what string, err error) error { + for _, kept := range []error{ErrNotOwner, ErrRelayUnavailable, ErrRelayRequestInvalid, ErrNoBulkSession, ErrProtocol} { + if errors.Is(err, kept) { + return fmt.Errorf("through replica %q: %w", owner, err) + } + } + return unreachablePeer(owner, fmt.Errorf("%s: %w", what, err)) } // handshake names the service of the worker on a stream and waits for the answer diff --git a/core/services/tunnel/dialer_test.go b/core/services/tunnel/dialer_test.go index d5a007623a91..3954b1ebe7a3 100644 --- a/core/services/tunnel/dialer_test.go +++ b/core/services/tunnel/dialer_test.go @@ -3,6 +3,7 @@ package tunnel_test import ( "bytes" "context" + "encoding/binary" "errors" "io" "net" @@ -26,6 +27,7 @@ type scriptedWorker struct { requests []workerRequest refuse error hold time.Duration + garbage bool } type workerRequest struct{ tag, target string } @@ -45,8 +47,12 @@ func (w *scriptedWorker) serve(sess *tunnel.Session) { } w.mu.Lock() w.requests = append(w.requests, workerRequest{tag, target}) - refuse, hold := w.refuse, w.hold + refuse, hold, garbage := w.refuse, w.hold, w.garbage w.mu.Unlock() + if garbage { + _, _ = st.Write(rawFrame("this is not a reply")) + return + } if hold > 0 { time.Sleep(hold) } @@ -178,6 +184,29 @@ var _ = Describe("The worker dialer", func() { Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) }) + It("reports a worker that says nothing for the whole backstop as a slow path and not as a missing route", func() { + hold(false) + worker.hold = 3 * time.Second + DeferCleanup(tunnel.SetHandshakeTimeout(300 * time.Millisecond)) + dialer := tunnel.NewWorkerDialer(tunnels, nil) + + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrTransport)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "a worker that is slow to answer is still held") + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("reports a worker that answers with something that is no reply as a protocol error and not as a missing route", func() { + hold(false) + worker.garbage = true + dialer := tunnel.NewWorkerDialer(tunnels, nil) + + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrProtocol)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse()) + Expect(tunnel.IsWorkerAnswer(err)).To(BeFalse()) + }) + It("reports a session that ended under a held entry as a route that does not exist now", func() { front, _ := hold(false) Expect(front.Close()).To(Succeed()) @@ -299,6 +328,30 @@ var _ = Describe("The worker dialer", func() { Eventually(wedged.farClosed, 2*time.Second).Should(BeTrue()) }) + It("reports an owner link that dies before the relay reply as an unreachable peer and not as a missing route", func() { + dialer := tunnel.NewWorkerDialer(tunnels, &scriptedPeers{onFrame: func(far net.Conn) { _ = far.Close() }}) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "a link between two frontends says nothing about the worker") + }) + + It("reports an owner that is silent for the whole backstop as an unreachable peer", func() { + DeferCleanup(tunnel.SetHandshakeTimeout(300 * time.Millisecond)) + dialer := tunnel.NewWorkerDialer(tunnels, &wedgedPeers{}) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse()) + }) + + It("reports an owner that answers with something that is no reply as a protocol error", func() { + dialer := tunnel.NewWorkerDialer(tunnels, &scriptedPeers{onFrame: func(far net.Conn) { + _, _ = far.Write(rawFrame("hello there")) + }}) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrProtocol)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse()) + }) + It("states the remaining budget of the caller in the relay request", func() { dialer := tunnel.NewWorkerDialer(tunnels, peers) budgeted, cancel := context.WithTimeout(ctx, 20*time.Second) @@ -369,6 +422,19 @@ var _ = Describe("The worker dialer", func() { }) }) + Describe("when the database cannot say which replica holds the tunnel", func() { + It("reports an infrastructure error and not a missing route", func() { + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + Expect(sqlDB.Close()).To(Succeed()) + + dialer := tunnel.NewWorkerDialer(tunnels, &fakePeers{}) + _, err = dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrInfrastructure)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse(), "a database that does not answer is no statement about a worker") + }) + }) + Describe("when no live replica holds the tunnel", func() { It("answers ErrNoRoute and never absence, for a worker with no connection row", func() { dialer := tunnel.NewWorkerDialer(tunnels, &fakePeers{}) @@ -500,3 +566,30 @@ func (w *wedgedPeers) farClosed() bool { defer w.mu.Unlock() return w.end } + +// rawFrame is a frame of the tunnel protocol: a two byte length and the payload. +func rawFrame(payload string) []byte { + buf := make([]byte, 2+len(payload)) + binary.BigEndian.PutUint16(buf[:2], uint16(len(payload))) + copy(buf[2:], payload) + return buf +} + +// scriptedPeers hands out a stream to an owner that reads the relay frame and +// then does what a spec says with its end of the stream. +type scriptedPeers struct{ onFrame func(far net.Conn) } + +func (p *scriptedPeers) Open(context.Context, string) (net.Conn, error) { + near, far := net.Pipe() + go func() { + var size [2]byte + if _, err := io.ReadFull(far, size[:]); err != nil { + return + } + if _, err := io.CopyN(io.Discard, far, int64(binary.BigEndian.Uint16(size[:]))); err != nil { + return + } + p.onFrame(far) + }() + return near, nil +} diff --git a/core/services/tunnel/export_test.go b/core/services/tunnel/export_test.go index c5d74c431280..be8087ded7f9 100644 --- a/core/services/tunnel/export_test.go +++ b/core/services/tunnel/export_test.go @@ -21,3 +21,14 @@ func SetClaimTimeout(d time.Duration) (restore func()) { func NewRelayWithTimeouts(tunnels *Registry, header, open time.Duration) *Relay { return newRelay(tunnels, header, open) } + +// RouteFailure exposes the classification of a failed dial to the specs. +func RouteFailure(nodeID string, cause error) error { return routeFailure(nodeID, cause) } + +// SetHandshakeTimeout changes the backstop of the handshake for one spec and +// returns a function that restores it. +func SetHandshakeTimeout(d time.Duration) (restore func()) { + old := dialHandshakeTimeout + dialHandshakeTimeout = d + return func() { dialHandshakeTimeout = old } +} diff --git a/core/services/tunnel/proto.go b/core/services/tunnel/proto.go index d53d543ff3e9..3fd379e70e48 100644 --- a/core/services/tunnel/proto.go +++ b/core/services/tunnel/proto.go @@ -223,7 +223,7 @@ func ReadStreamReply(r io.Reader) error { } rest, ok := strings.CutPrefix(payload, replyPrefixRefused) if !ok { - return fmt.Errorf("reading a tunnel stream reply: unrecognised reply %q", payload) + return fmt.Errorf("reading a tunnel stream reply: %w: %q", ErrProtocol, payload) } code, text, _ := strings.Cut(rest, streamRequestSeparator) for _, r := range streamRefusals { @@ -234,7 +234,7 @@ func ReadStreamReply(r io.Reader) error { // A code from a newer worker. It is not mapped onto the nearest known code, // so that a frontend does not retry for ever against a refusal with another // meaning, and so that IsWorkerAnswer is false and nothing is reaped. - return fmt.Errorf("tunnel stream refused with unrecognised code %q: %s", code, text) + return fmt.Errorf("tunnel stream refused with unrecognised code %q: %w: %s", code, ErrProtocol, text) } // truncateRunes cuts s to at most limit bytes on a rune boundary. A plain slice @@ -279,7 +279,7 @@ func readFrame(r io.Reader) (string, error) { } n := binary.BigEndian.Uint16(size[:]) if int(n) > maxFrame { - return "", fmt.Errorf("tunnel frame declares %d bytes, over the %d-byte limit", n, maxFrame) + return "", fmt.Errorf("%w: tunnel frame declares %d bytes, over the %d-byte limit", ErrProtocol, n, maxFrame) } if n == 0 { return "", nil diff --git a/core/services/tunnel/relay.go b/core/services/tunnel/relay.go index dd8df8d6e42a..f1da017d9e11 100644 --- a/core/services/tunnel/relay.go +++ b/core/services/tunnel/relay.go @@ -189,7 +189,7 @@ func ReadRelayReply(r io.Reader) error { } rest, ok := strings.CutPrefix(payload, replyPrefixRefused) if !ok { - return fmt.Errorf("reading a relay reply: unrecognised reply %q", payload) + return fmt.Errorf("reading a relay reply: %w: %q", ErrProtocol, payload) } code, text, _ := strings.Cut(rest, streamRequestSeparator) switch code { @@ -208,7 +208,7 @@ func ReadRelayReply(r io.Reader) error { // A code of a newer replica. It is not mapped onto the nearest known code, so // that a caller does not retry for ever against a refusal that means something // else. - return fmt.Errorf("relay stream refused with unrecognised code %q: %s", code, text) + return fmt.Errorf("relay stream refused with unrecognised code %q: %w: %s", code, ErrProtocol, text) } // maxRelayBudgetMillis is the largest budget a peer may declare. A day is far diff --git a/core/services/tunnel/route.go b/core/services/tunnel/route.go index 6164cf06d904..56de3bedef63 100644 --- a/core/services/tunnel/route.go +++ b/core/services/tunnel/route.go @@ -4,6 +4,11 @@ import ( "context" "errors" "fmt" + "io" + "net" + "os" + + "github.com/libp2p/go-yamux/v5" "github.com/mudler/LocalAI/core/services/cluster" ) @@ -23,6 +28,8 @@ import ( // routeFailure describes. // // It is not carried by: +// - a path that is slow (ErrTransport), a reply that breaks the protocol +// (ErrProtocol) or a database that did not answer (ErrInfrastructure); // - a refusal that the worker wrote, because a worker that answers is present; // - a peer replica that did not answer (ErrPeerUnreachable), because that is a // fact about the link between two frontends and a scheduler that demoted a @@ -37,6 +44,23 @@ var ErrNoRoute = errors.New("tunnel: no route from this replica to that worker") // ErrPeerUnreachable, because no peer was dialled. var ErrNoRelayPath = errors.New("tunnel: this replica cannot relay to the owner of that worker") +// ErrTransport reports that the tunnel path to a worker did not carry a request +// in time: a socket deadline ran out, a keepalive was missed, a worker took the +// request and said nothing. It is a fact about how fast the path is, not about +// whether a route exists. A scheduler that demoted a worker on it would take +// capacity away from a worker that is slow to answer. +var ErrTransport = errors.New("tunnel: the path to that worker did not carry the request in time") + +// ErrProtocol reports a reply on a stream that is not part of the protocol of +// the tunnel: a worker or a relay that answered with something else. It is a +// defect of the other end, and not a missing route. +var ErrProtocol = errors.New("tunnel: unexpected reply on a stream") + +// ErrInfrastructure reports that the shared database, which holds the owner of +// each tunnel, could not answer. It says nothing about any worker: the worker +// may be well, and so may the replica that holds it. +var ErrInfrastructure = errors.New("tunnel: the cluster database could not answer") + // noRouteError reports a worker that this replica cannot route to, keeping the // cause in its message and out of its unwrap chain. // @@ -69,19 +93,57 @@ func isAbsenceClaim(err error) bool { // - An absence claim does not reach the caller. It becomes ErrNoRoute with the // cause in the message only. // - A failure of the budget of the caller, of the link to a peer replica, of the -// pool, or of the bulk lane alone is returned as itself, wrapped, with no -// ErrNoRoute. None of them is a statement about whether the worker has a -// route, and ErrNoRoute makes a scheduler demote the worker. -// - Everything else is a route that does not exist now: ErrNoRoute on top, and -// the cause below it. +// pool, of the bulk lane alone, of the database, or a reply that breaks the +// protocol is returned as itself, wrapped, with no ErrNoRoute. None of them is +// a statement about whether the worker has a route, and ErrNoRoute makes a +// scheduler demote the worker. +// - A timeout of a socket or of the multiplexer is ErrTransport. A path that is +// slow is not a path that does not exist. +// - ErrNoRoute is for the routing facts: nothing holds the tunnel of the node, +// this replica cannot reach the one that does, the owner could not open a +// stream, the worker could not serve one, or the session ended under the +// request. The cause stays below it. +// - Anything else is returned as a plain wrapped error. A failure that nobody +// classified must not demote a worker. func routeFailure(nodeID string, cause error) error { switch { case isAbsenceClaim(cause): return &noRouteError{nodeID: nodeID, cause: cause} case errors.Is(cause, context.DeadlineExceeded), errors.Is(cause, context.Canceled), errors.Is(cause, ErrPeerUnreachable), errors.Is(cause, ErrPoolClosed), - errors.Is(cause, ErrNoBulkSession): + errors.Is(cause, ErrNoBulkSession), errors.Is(cause, ErrInfrastructure), + errors.Is(cause, ErrProtocol), errors.Is(cause, ErrTransport): return fmt.Errorf("reaching node %q: %w", nodeID, cause) + case isTimeout(cause): + return fmt.Errorf("reaching node %q: %w: %w", nodeID, ErrTransport, cause) + case isRoutingFact(cause): + return fmt.Errorf("reaching node %q: %w: %w", nodeID, ErrNoRoute, cause) + } + return fmt.Errorf("reaching node %q: %w", nodeID, cause) +} + +// isTimeout reports whether a deadline of a socket or of the multiplexer ended +// the attempt. +func isTimeout(err error) bool { + if errors.Is(err, os.ErrDeadlineExceeded) || errors.Is(err, yamux.ErrTimeout) { + return true + } + var netErr net.Error + return errors.As(err, &netErr) && netErr.Timeout() +} + +// isRoutingFact reports whether err says that the path to the worker is not +// there: not held, not reachable through a relay, or ended under the request. +func isRoutingFact(err error) bool { + for _, fact := range []error{ + ErrNotOwner, ErrNoRelayPath, ErrRelayUnavailable, ErrRelayRequestInvalid, + ErrStreamNotServed, + yamux.ErrSessionShutdown, yamux.ErrStreamClosed, yamux.ErrStreamReset, yamux.ErrRemoteGoAway, + io.EOF, io.ErrUnexpectedEOF, io.ErrClosedPipe, net.ErrClosed, + } { + if errors.Is(err, fact) { + return true + } } - return fmt.Errorf("reaching node %q: %w: %w", nodeID, ErrNoRoute, cause) + return false } diff --git a/core/services/tunnel/route_test.go b/core/services/tunnel/route_test.go new file mode 100644 index 000000000000..1ed10355df87 --- /dev/null +++ b/core/services/tunnel/route_test.go @@ -0,0 +1,60 @@ +package tunnel_test + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "os" + + "github.com/libp2p/go-yamux/v5" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// netTimeout is a timeout of a socket that is not os.ErrDeadlineExceeded. +type netTimeout struct{} + +func (netTimeout) Error() string { return "i/o timeout" } +func (netTimeout) Timeout() bool { return true } +func (netTimeout) Temporary() bool { return true } + +var _ net.Error = netTimeout{} + +// The failure of a dial reaches a scheduler, and the scheduler demotes a worker +// on ErrNoRoute. ErrNoRoute therefore says that no route exists, and a path that +// is slow, a peer that is gone, a database that does not answer and a reply that +// makes no sense are other conditions. +var _ = DescribeTable("The classification of a failed dial", + func(cause error, want error, wantNoRoute bool) { + got := tunnel.RouteFailure("w1", cause) + if want != nil { + Expect(errors.Is(got, want)).To(BeTrue(), "%v", got) + } + Expect(errors.Is(got, tunnel.ErrNoRoute)).To(Equal(wantNoRoute), "%v", got) + Expect(errors.Is(got, cluster.ErrNoConnection)).To(BeFalse(), "absence must not reach the caller") + Expect(errors.Is(got, cluster.ErrInstanceNotFound)).To(BeFalse()) + }, + Entry("a socket deadline that ran out", fmt.Errorf("reading a tunnel stream reply: %w", os.ErrDeadlineExceeded), tunnel.ErrTransport, false), + Entry("a timeout of the network", fmt.Errorf("writing: %w", netTimeout{}), tunnel.ErrTransport, false), + Entry("a timeout of the multiplexer", fmt.Errorf("opening a stream: %w", yamux.ErrTimeout), tunnel.ErrTransport, false), + Entry("a reply that is no part of the protocol", fmt.Errorf("%w: %q", tunnel.ErrProtocol, "garbage"), tunnel.ErrProtocol, false), + Entry("a database that did not answer", fmt.Errorf("%w: connection refused", tunnel.ErrInfrastructure), tunnel.ErrInfrastructure, false), + Entry("a peer that did not answer", fmt.Errorf("through replica: %w", tunnel.ErrPeerUnreachable), tunnel.ErrPeerUnreachable, false), + Entry("the budget of the caller", context.DeadlineExceeded, context.DeadlineExceeded, false), + Entry("a cancelled caller", context.Canceled, context.Canceled, false), + Entry("a bulk lane that is down", tunnel.ErrNoBulkSession, tunnel.ErrNoBulkSession, false), + Entry("a node that no replica holds", cluster.ErrNoConnection, nil, true), + Entry("a replica that is not registered", cluster.ErrInstanceNotFound, nil, true), + Entry("a tunnel that is not held here", tunnel.ErrNotOwner, tunnel.ErrNotOwner, true), + Entry("a replica that cannot relay", tunnel.ErrNoRelayPath, tunnel.ErrNoRelayPath, true), + Entry("an owner that could not open a stream", tunnel.ErrRelayUnavailable, tunnel.ErrRelayUnavailable, true), + Entry("a worker that could not serve the stream", tunnel.ErrStreamNotServed, tunnel.ErrStreamNotServed, true), + Entry("a session that ended under the request", yamux.ErrSessionShutdown, yamux.ErrSessionShutdown, true), + Entry("a stream that was closed under the request", fmt.Errorf("reading: %w", io.EOF), io.EOF, true), +) From b57c3552e797b59a0c834083a0a3bdecdacc9134 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 07:53:16 +0000 Subject: [PATCH 093/162] fix(tunnel): look the owner up again when the owner no longer holds the tunnel A worker that moves to another replica between the lookup and the relay made the old owner refuse with ErrNotOwner, and the dial ended as a missing route. The dialer now looks the owner up once more, within the budget of the caller, and uses the local tunnel when the worker moved to this replica. The second answer is final, so a stale row cannot turn into a loop. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/tunnel/dialer.go | 4 ++ core/services/tunnel/dialer_test.go | 66 +++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/core/services/tunnel/dialer.go b/core/services/tunnel/dialer.go index 8a9462ec755a..e54f90d2ed7e 100644 --- a/core/services/tunnel/dialer.go +++ b/core/services/tunnel/dialer.go @@ -127,6 +127,10 @@ func (d *WorkerDialer) GRPCDialerFor(nodeID string) func(ctx context.Context, ad // answer is final: a third attempt would turn a stale row into a loop. func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string, o dialOptions) (net.Conn, error) { stream, err := d.relayOnce(ctx, nodeID, o, false) + if err != nil && errors.Is(err, ErrNotOwner) && errors.Is(err, errOwnerRefused) && callerRanOut(ctx) == nil { + xlog.Debug("the owner of a tunnel no longer holds it, looking the owner up again", "node", nodeID) + stream, err = d.relayOnce(ctx, nodeID, o, true) + } if err != nil { return nil, routeFailure(nodeID, err) } diff --git a/core/services/tunnel/dialer_test.go b/core/services/tunnel/dialer_test.go index 3954b1ebe7a3..1ffd26e72ec0 100644 --- a/core/services/tunnel/dialer_test.go +++ b/core/services/tunnel/dialer_test.go @@ -385,6 +385,38 @@ var _ = Describe("The worker dialer", func() { Expect(errors.Is(err, cluster.ErrNoConnection)).To(BeFalse(), "absence must not reach the caller") }) + It("looks the owner up again once when the owner says it no longer holds the tunnel", func() { + // The worker moved from replica-b to replica-c after the lookup of the + // dialer. The first owner refuses, and the second one serves. + Expect(clusterR.Register(ctx, "replica-c", "test", 0, "")).To(Succeed()) + third := tunnel.NewRegistry(clusterR, "replica-c") + moving := &movingPeers{ + relays: map[string]*tunnel.Relay{"replica-b": tunnel.NewRelay(other), "replica-c": tunnel.NewRelay(third)}, + moveOnFirst: func() { + Expect(other.Disconnect("w1")).To(BeTrue()) + front, back := sessionPair(tunnel.LaneInference) + worker.serve(back) + _, err := third.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + }, + } + dialer := tunnel.NewWorkerDialer(tunnels, moving) + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + roundTrip(conn, "after the move") + Expect(moving.opened()).To(Equal([]string{"replica-b", "replica-c"})) + }) + + It("looks the owner up only once more, and then reports the routing fact", func() { + stale := &movingPeers{relays: map[string]*tunnel.Relay{"replica-b": tunnel.NewRelay(tunnel.NewRegistry(clusterR, "replica-b"))}} + dialer := tunnel.NewWorkerDialer(tunnels, stale) + _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrNotOwner)).To(BeTrue()) + Expect(stale.opened()).To(HaveLen(2), "one lookup and one more") + }) + It("reports an unreachable peer as unreachable, with no ErrNoRoute and no absence", func() { dialer := tunnel.NewWorkerDialer(tunnels, &downPeers{err: tunnelUnreachable()}) _, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "x:1") @@ -593,3 +625,37 @@ func (p *scriptedPeers) Open(context.Context, string) (net.Conn, error) { }() return near, nil } + +// movingPeers relays to the replica that a dial names. The first open can move +// the worker to another replica before it answers, which is what a worker that +// reconnects to a different replica does between a lookup and a dial. +type movingPeers struct { + relays map[string]*tunnel.Relay + moveOnFirst func() + + mu sync.Mutex + peers []string +} + +func (m *movingPeers) Open(_ context.Context, peerID string) (net.Conn, error) { + m.mu.Lock() + first := len(m.peers) == 0 + m.peers = append(m.peers, peerID) + m.mu.Unlock() + if first && m.moveOnFirst != nil { + m.moveOnFirst() + } + relay, ok := m.relays[peerID] + if !ok { + return nil, errors.New("no relay behind this fake") + } + near, far := net.Pipe() + go relay.Stream("replica-a", far) + return near, nil +} + +func (m *movingPeers) opened() []string { + m.mu.Lock() + defer m.mu.Unlock() + return append([]string(nil), m.peers...) +} From 372bca0851344bdca2d3b8e7db0f9136ad4d14b5 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 07:54:44 +0000 Subject: [PATCH 094/162] fix(nodes): serve the control verbs of a worker only to a caller that may use them checkBearerToken returns true for an empty token, and the control plane used it. A worker on the tunnel carrier that was started with an address binds every interface, so backend.install (from any URI), node.stop, model.delete and model.op were open to anyone who could reach the port when the deployment had no registration token. The control routes now go through controlGate. With a token, the bearer check decides. Without one, only a caller on the loopback address is served, which is the stream that the tunnel opens on the worker, and every other caller is refused for every verb. The file routes keep their behaviour on an empty token. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 6 +++ .../nodes/file_transfer_control_test.go | 46 ++++++++++++++++ core/services/nodes/file_transfer_server.go | 52 ++++++++++++++++--- 3 files changed, 97 insertions(+), 7 deletions(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 770a310f97ce..fe2452680121 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -343,6 +343,12 @@ back to the lane of model calls. The file stager asks for it. ## Rules a carrier must keep +- The control verbs of a worker are never served unauthenticated to the network. + `nodes.controlGate` checks the bearer token, and with no token it serves only a + caller on the loopback address, which is the stream that the tunnel opens on + the worker. The file routes keep their older behaviour on an empty token. An + agent worker binds loopback only. + - Payload: a broadcast above `messaging.MaxBroadcastBytes` is refused with `messaging.ErrPayloadTooLarge` by every carrier. - Subjects: one closed set of roots, the `broadcastRoots` and `controlRoots` diff --git a/core/services/nodes/file_transfer_control_test.go b/core/services/nodes/file_transfer_control_test.go index f5852200e908..685170d29a2f 100644 --- a/core/services/nodes/file_transfer_control_test.go +++ b/core/services/nodes/file_transfer_control_test.go @@ -5,6 +5,7 @@ import ( "io" "net" "net/http" + "net/http/httptest" "strings" "time" @@ -105,4 +106,49 @@ var _ = Describe("The control plane on the file transfer server", func() { defer func() { _ = resp.Body.Close() }() Expect(resp.StatusCode).To(Equal(http.StatusOK)) }) + + Describe("on a worker that has no registration token", func() { + It("refuses every control verb to a caller that is not on this host", func() { + var ran int + gate := controlGate("", http.HandlerFunc(func(http.ResponseWriter, *http.Request) { ran++ })) + for _, verb := range workerctl.AllVerbs() { + req := httptest.NewRequest(http.MethodPost, workerctl.PathOf(verb), strings.NewReader("{}")) + req.RemoteAddr = "203.0.113.7:51234" + rec := httptest.NewRecorder() + gate.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusForbidden), verb) + } + Expect(ran).To(BeZero(), "no verb may run for a caller from the network") + }) + + It("refuses a caller whose address cannot be read", func() { + gate := controlGate("", http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + req := httptest.NewRequest(http.MethodPost, workerctl.PathOf(workerctl.VerbNodeStop), nil) + req.RemoteAddr = "not an address" + rec := httptest.NewRecorder() + gate.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusForbidden)) + }) + + It("serves the stream that the tunnel opens on the loopback address", func() { + control := http.NewServeMux() + control.HandleFunc(workerctl.PathOf(workerctl.VerbBackendList), func(w http.ResponseWriter, _ *http.Request) { + _, _ = fmt.Fprint(w, "listed") + }) + base := start("", control) + status, body := do(base+workerctl.PathOf(workerctl.VerbBackendList), "") + Expect(status).To(Equal(http.StatusOK)) + Expect(body).To(Equal("listed")) + }) + + It("lets the bearer check decide, and not the address, when there is a token", func() { + gate := controlGate("secret", http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNoContent) })) + req := httptest.NewRequest(http.MethodPost, workerctl.PathOf(workerctl.VerbNodeStop), nil) + req.RemoteAddr = "203.0.113.7:51234" + req.Header.Set("Authorization", "Bearer secret") + rec := httptest.NewRecorder() + gate.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusNoContent)) + }) + }) }) diff --git a/core/services/nodes/file_transfer_server.go b/core/services/nodes/file_transfer_server.go index 97fce01921f1..bce11ee4d447 100644 --- a/core/services/nodes/file_transfer_server.go +++ b/core/services/nodes/file_transfer_server.go @@ -109,6 +109,11 @@ func startFileTransferServerWithControl(lis net.Listener, stagingDir, modelsDir, xlog.Warn("HTTP file transfer server starting WITHOUT a registration token — read/write to models/staging/data is unauthenticated for anyone who can reach this port; set LOCALAI_REGISTRATION_TOKEN") } + if control != nil && token == "" { + xlog.Warn("The control plane of this worker has no registration token, so it serves only callers on this host. " + + "A frontend reaches it through the tunnel; set LOCALAI_REGISTRATION_TOKEN to serve it to a caller on the network") + } + mux := http.NewServeMux() // PUT /v1/files/{key} — upload file @@ -177,13 +182,7 @@ func startFileTransferServerWithControl(lis net.Listener, stagingDir, modelsDir, // a tunnel. The bearer check is the one of the file routes: a verb can stop a // node, so it is never served without it. if control != nil { - mux.HandleFunc(workerctl.Prefix, func(w http.ResponseWriter, r *http.Request) { - if !checkBearerToken(r, token) { - http.Error(w, "unauthorized", http.StatusUnauthorized) - return - } - control.ServeHTTP(w, r) - }) + mux.Handle(workerctl.Prefix, controlGate(token, control)) } // Backend log endpoints (only registered when a log store is provided) @@ -1081,6 +1080,45 @@ func resolveKeyToDir(key, stagingDir, modelsDir, dataDir string) (targetDir, rel return } +// controlGate puts the check of the control plane in front of next. +// +// The file routes have always fail open on an empty token, and a worker on NATS +// keeps that. The control verbs are worse: they install a backend from any URI, +// stop the node and delete models. A worker on the tunnel carrier that was given +// an address binds every interface, so a control plane that failed open would be +// open to the network. With a token, the bearer check decides, as for the file +// routes. Without one, only a caller on the loopback address is served, which is +// the stream that the tunnel opens on the worker, and any other caller is +// refused for every verb. +func controlGate(token string, next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if token == "" { + if !fromLoopback(r) { + http.Error(w, "the control plane needs a registration token for a caller that is not on this host", http.StatusForbidden) + return + } + next.ServeHTTP(w, r) + return + } + if !checkBearerToken(r, token) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + next.ServeHTTP(w, r) + }) +} + +// fromLoopback reports whether the peer of the request is on the loopback +// interface of this host. +func fromLoopback(r *http.Request) bool { + host, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + return false + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + // checkBearerToken validates a Bearer token from the Authorization header // using constant-time comparison. Returns true if valid or if expectedToken is empty. func checkBearerToken(r *http.Request, expectedToken string) bool { From 6912c5eaed59e3a76786e26ea49aaf17c6dec7d4 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 07:56:29 +0000 Subject: [PATCH 095/162] feat(tunnel): let the peer link run over TLS and warn when it does not The peer link was always plain ws, so the credential of a replica and every request that a replica relays, prompts included, crossed it in clear text. PeerPool takes a WithPeerTLS option and dials wss. The flags --peer-tls and --peer-tls-ca (LOCALAI_PEER_TLS, LOCALAI_PEER_TLS_CA) turn it on and name a private CA. The peers must sit behind TLS at the address that they publish, for example behind a reverse proxy. A replica that publishes an address that is not on its own host warns at start-up while the link is plain. The flags are documented next to --peer-address. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 4 +- core/application/distributed.go | 57 +++++++++++++++++- core/application/peer_tls_test.go | 53 ++++++++++++++++ core/cli/run.go | 5 ++ core/config/distributed_config.go | 18 +++++- core/services/tunnel/peerlink.go | 38 ++++++++++-- core/services/tunnel/peerpool_test.go | 73 +++++++++++++++++++++++ docs/content/features/distributed-mode.md | 4 +- 8 files changed, 241 insertions(+), 11 deletions(-) create mode 100644 core/application/peer_tls_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index fe2452680121..4e885e110fc0 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -317,7 +317,9 @@ any other replica relays through the owner over a link that the pool of peers (`tunnel.PeerSessions`, `tunnel.Relay`). A replica publishes its address and the hash of its own peer credential in its instances row. The credential of the replica is the only secret that opens a peer link; the registration token opens -none. A relayed stream carries two request frames, the relay frame (worker, lane, +none. The link is plain `ws` unless `--peer-tls` is set (`tunnel.WithPeerTLS`), +and a replica that publishes an address that is not on its host warns at +start-up while it is plain. A relayed stream carries two request frames, the relay frame (worker, lane, remaining budget of the caller) and then the frame of the tunnel. The relay never relays onward. diff --git a/core/application/distributed.go b/core/application/distributed.go index aeb0a3386cc4..1f368ff92605 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -2,10 +2,13 @@ package application import ( "context" + "crypto/tls" + "crypto/x509" "encoding/json" "fmt" "io" "net" + "os" "strconv" "strings" "sync" @@ -603,7 +606,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // itself. Only the hash is published, in the same row as the address that // the peers dial. peerCred := cluster.NewPeerCredential() - membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID, peerAddress(cfg), peerCred) + advertised := peerAddress(cfg) + membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID, advertised, peerCred) if err != nil { return nil, err } @@ -619,7 +623,11 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // If a peer sweeps this replica while it stalls, the claims of the tunnels // that it still holds are written again when it registers again. membership.SetReclaimer(tunnels) - peerPool := tunnel.NewPeerPool(cfg.Distributed.InstanceID, peerCred, clusterReg) + peerOpts, err := peerPoolOptions(cfg, advertised) + if err != nil { + return nil, err + } + peerPool := tunnel.NewPeerPool(cfg.Distributed.InstanceID, peerCred, clusterReg, peerOpts...) peerSessions := tunnel.NewPeerSessions(tunnel.NewRelay(tunnels).Stream) success = true @@ -682,6 +690,51 @@ func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } +// peerPoolOptions reads the TLS settings of the peer link. It also warns when the +// link would cross a network in clear text: the credential of this replica and +// every relayed request, prompts included, go over it. +func peerPoolOptions(cfg *config.ApplicationConfig, advertised string) ([]tunnel.PeerPoolOption, error) { + d := cfg.Distributed + if !d.PeerTLS && d.PeerTLSCA == "" { + if addr := advertised; clearTextOnNetwork(addr) { + xlog.Warn("The peer link is clear text and this replica publishes an address that is not on this host, so the credential of this replica and the requests that other replicas relay cross the network unencrypted. "+ + "Put the replicas behind TLS and set LOCALAI_PEER_TLS (and LOCALAI_PEER_TLS_CA for a private CA)", "address", addr) + } + return nil, nil + } + tlsCfg := &tls.Config{MinVersion: tls.VersionTLS12} + if d.PeerTLSCA != "" { + pem, err := os.ReadFile(d.PeerTLSCA) + if err != nil { + return nil, fmt.Errorf("reading the CA file of the peer link: %w", err) + } + roots := x509.NewCertPool() + if !roots.AppendCertsFromPEM(pem) { + return nil, fmt.Errorf("the CA file of the peer link, %s, holds no certificate", d.PeerTLSCA) + } + tlsCfg.RootCAs = roots + } + return []tunnel.PeerPoolOption{tunnel.WithPeerTLS(tlsCfg)}, nil +} + +// clearTextOnNetwork reports whether an address published for the peer link +// names a host other than this one. An empty address is not published, so +// nothing is dialled. +func clearTextOnNetwork(addr string) bool { + if addr == "" { + return false + } + host, _, err := net.SplitHostPort(addr) + if err != nil { + host = addr + } + if strings.EqualFold(host, "localhost") { + return false + } + ip := net.ParseIP(host) + return ip == nil || !ip.IsLoopback() +} + // peerAddress returns the address at which the other replicas dial this one. An // address that the operator configured is checked and used as it is. Without // one, it comes from the route to the database, which all replicas share. diff --git a/core/application/peer_tls_test.go b/core/application/peer_tls_test.go new file mode 100644 index 000000000000..7185e332f788 --- /dev/null +++ b/core/application/peer_tls_test.go @@ -0,0 +1,53 @@ +package application + +import ( + "os" + "path/filepath" + + "github.com/mudler/LocalAI/core/config" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("The TLS settings of the peer link", func() { + DescribeTable("warns about clear text only for an address that is not on this host", + func(addr string, want bool) { + Expect(clearTextOnNetwork(addr)).To(Equal(want)) + }, + Entry("nothing published", "", false), + Entry("loopback", "127.0.0.1:8080", false), + Entry("loopback v6", "[::1]:8080", false), + Entry("localhost", "localhost:8080", false), + Entry("a private address", "10.1.2.3:8080", true), + Entry("a public address", "203.0.113.9:443", true), + Entry("a name", "frontend-2.internal:8080", true), + ) + + It("adds no option and keeps the link plain when TLS is off", func() { + opts, err := peerPoolOptions(&config.ApplicationConfig{}, "10.1.2.3:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(opts).To(BeEmpty()) + }) + + It("adds the TLS option when it is on, with the system roots", func() { + cfg := &config.ApplicationConfig{} + cfg.Distributed.PeerTLS = true + opts, err := peerPoolOptions(cfg, "10.1.2.3:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(opts).To(HaveLen(1)) + }) + + It("refuses a CA file that holds no certificate, and one that is missing", func() { + cfg := &config.ApplicationConfig{} + cfg.Distributed.PeerTLSCA = filepath.Join(GinkgoT().TempDir(), "missing.pem") + _, err := peerPoolOptions(cfg, "") + Expect(err).To(HaveOccurred()) + + empty := filepath.Join(GinkgoT().TempDir(), "empty.pem") + Expect(os.WriteFile(empty, []byte("not a certificate"), 0o600)).To(Succeed()) + cfg.Distributed.PeerTLSCA = empty + _, err = peerPoolOptions(cfg, "") + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/cli/run.go b/core/cli/run.go index e64cf2321865..129d1d7bdcc2 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -166,6 +166,8 @@ type RunCMD struct { Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"` InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"` PeerAddress string `env:"LOCALAI_PEER_ADDRESS" help:"host:port at which the other frontends dial this one to reach a worker tunnel it holds (found from the route to the database if empty)" group:"distributed"` + PeerTLS bool `env:"LOCALAI_PEER_TLS" default:"false" help:"dial the other frontends over wss; they must sit behind TLS at the address they publish. Without it the peer link, which carries the credential of this replica and relayed requests, is clear text" group:"distributed"` + PeerTLSCA string `env:"LOCALAI_PEER_TLS_CA" type:"existingfile" help:"PEM file with the CA that signs the certificate of the other frontends (system roots if empty); use with --peer-tls" group:"distributed"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` @@ -367,6 +369,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.PeerAddress != "" { opts = append(opts, config.WithDistributedPeerAddress(r.PeerAddress)) } + if r.PeerTLS || r.PeerTLSCA != "" { + opts = append(opts, config.WithDistributedPeerTLS(r.PeerTLSCA)) + } if r.NatsURL != "" { opts = append(opts, config.WithNatsURL(r.NatsURL)) } diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index b02d666a095b..22f8b1488c98 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -18,7 +18,14 @@ type DistributedConfig struct { // PeerAddress is the host and port at which the other frontends dial this // one, to reach a worker tunnel that this one holds. Empty means that the // address is found from the route to the database. LOCALAI_PEER_ADDRESS. - PeerAddress string + PeerAddress string + // PeerTLS makes this frontend dial the other frontends over wss. They must sit + // behind TLS at the address that they publish. LOCALAI_PEER_TLS. + PeerTLS bool + // PeerTLSCA is a PEM file with the certificate authority that signs the + // certificate of the other frontends. Empty uses the system roots. + // LOCALAI_PEER_TLS_CA. + PeerTLSCA string NatsURL string // --nats-url / LOCALAI_NATS_URL StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) @@ -208,6 +215,15 @@ func WithDistributedPeerAddress(addr string) AppOption { } } +// WithDistributedPeerTLS makes this frontend dial its peers over wss, trusting +// the CA in caFile, or the system roots when caFile is empty. +func WithDistributedPeerTLS(caFile string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.PeerTLS = true + o.Distributed.PeerTLSCA = caFile + } +} + func WithNatsURL(url string) AppOption { return func(o *ApplicationConfig) { o.Distributed.NatsURL = url diff --git a/core/services/tunnel/peerlink.go b/core/services/tunnel/peerlink.go index 4756f6b512b2..fe829b15653c 100644 --- a/core/services/tunnel/peerlink.go +++ b/core/services/tunnel/peerlink.go @@ -2,6 +2,7 @@ package tunnel import ( "context" + "crypto/tls" "errors" "fmt" "net" @@ -108,6 +109,8 @@ type PeerPool struct { reg *cluster.Registry dialer *websocket.Dialer + // secure says that peers are dialled over wss. + secure bool mu sync.Mutex links map[string]*peerLink @@ -166,8 +169,8 @@ func (l *peerLink) setSession(s *Session) { // The pool authenticates with the credential of the replica and with nothing // else. The registration token of the deployment opens no peer link, and an empty // token therefore fails no peer dial. -func NewPeerPool(selfID string, cred PeerCredential, reg *cluster.Registry) *PeerPool { - return &PeerPool{ +func NewPeerPool(selfID string, cred PeerCredential, reg *cluster.Registry, opts ...PeerPoolOption) *PeerPool { + p := &PeerPool{ selfID: selfID, cred: cred, reg: reg, @@ -176,6 +179,25 @@ func NewPeerPool(selfID string, cred PeerCredential, reg *cluster.Registry) *Pee dialer: NewDialer(peerLinkHandshakeTimeout), links: map[string]*peerLink{}, } + for _, opt := range opts { + opt(p) + } + return p +} + +// PeerPoolOption changes how a PeerPool dials. +type PeerPoolOption func(*PeerPool) + +// WithPeerTLS makes the pool dial its peers over wss and not ws, with the given +// client configuration. The peer must then sit behind TLS at the address that it +// publishes, for example behind the reverse proxy of the deployment. Without it +// the credential of this replica and every relayed request cross the link in +// clear text. +func WithPeerTLS(cfg *tls.Config) PeerPoolOption { + return func(p *PeerPool) { + p.dialer.TLSClientConfig = cfg.Clone() + p.secure = true + } } // Open returns a stream to peerID, and dials and caches the session on first use. @@ -280,11 +302,15 @@ func (p *PeerPool) dial(ctx context.Context, peerID string) (*Session, error) { return nil, unreachablePeer(peerID, errors.New("peer has no advertised address")) } + scheme := "ws" + if p.secure { + scheme = "wss" + } endpoint := url.URL{ - // Plain ws: the link is authenticated by the credential of the replica and - // not by the transport. A deployment that wants the link encrypted puts the - // replicas behind TLS. - Scheme: "ws", + // Plain ws unless WithPeerTLS was given: the link is authenticated by the + // credential of the replica, and the transport is encrypted only when the + // operator asked for it. + Scheme: scheme, Host: inst.AdvertisedAddr, Path: PeerPath, RawQuery: url.Values{"id": []string{p.selfID}}.Encode(), diff --git a/core/services/tunnel/peerpool_test.go b/core/services/tunnel/peerpool_test.go index e85f5bd1b7ad..8822f6e620f7 100644 --- a/core/services/tunnel/peerpool_test.go +++ b/core/services/tunnel/peerpool_test.go @@ -2,8 +2,12 @@ package tunnel_test import ( "context" + "crypto/tls" + "crypto/x509" "errors" "net" + "net/http" + "net/http/httptest" "sync" "time" @@ -77,4 +81,73 @@ var _ = Describe("The peer pool", func() { pool.Close() Eventually(first, 15*time.Second).Should(Receive()) }) + + Describe("over TLS", func() { + // peerServer is a replica that accepts a peer link at PeerPath, behind TLS. + peerServer := func() *httptest.Server { + GinkgoHelper() + upgrader := tunnel.NewUpgrader() + sessions := tunnel.NewPeerSessions(func(_ string, stream net.Conn) { _ = stream.Close() }) + DeferCleanup(sessions.Close) + mux := http.NewServeMux() + mux.HandleFunc(tunnel.PeerPath, func(w http.ResponseWriter, r *http.Request) { + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := tunnel.PeerServerSession(ws) + if err != nil { + _ = ws.Close() + return + } + sessions.Accept(r.URL.Query().Get("id"), sess) + }) + srv := httptest.NewTLSServer(mux) + DeferCleanup(srv.Close) + return srv + } + + poolFor := func(srv *httptest.Server, opts ...tunnel.PeerPoolOption) *tunnel.PeerPool { + GinkgoHelper() + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg := cluster.NewRegistry(db) + Expect(reg.RegisterPeer(ctx, "tls-peer", "test", 0, "", srv.Listener.Addr().String(), cluster.NewPeerCredential().Hash())).To(Succeed()) + pool := tunnel.NewPeerPool("replica-a", cluster.NewPeerCredential(), reg, opts...) + DeferCleanup(pool.Close) + return pool + } + + It("opens a stream to a peer that sits behind TLS, when it trusts the certificate", func() { + srv := peerServer() + roots := x509.NewCertPool() + roots.AddCert(srv.Certificate()) + pool := poolFor(srv, tunnel.WithPeerTLS(&tls.Config{RootCAs: roots, MinVersion: tls.VersionTLS12})) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + st, err := pool.Open(ctx, "tls-peer") + Expect(err).ToNot(HaveOccurred()) + _ = st.Close() + }) + + It("does not reach a peer behind TLS with a plain link, and says that the peer is unreachable", func() { + srv := peerServer() + pool := poolFor(srv) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, err := pool.Open(ctx, "tls-peer") + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeTrue(), "%v", err) + }) + + It("refuses a peer whose certificate it does not trust", func() { + srv := peerServer() + pool := poolFor(srv, tunnel.WithPeerTLS(&tls.Config{RootCAs: x509.NewCertPool(), MinVersion: tls.VersionTLS12})) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, err := pool.Open(ctx, "tls-peer") + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeTrue(), "%v", err) + }) + }) }) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index b29579c05a78..29dbff82edb8 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -63,7 +63,9 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These |------|---------|---------|-------------| | `--distributed` | `LOCALAI_DISTRIBUTED` | `false` | Enable distributed mode | | `--instance-id` | `LOCALAI_INSTANCE_ID` | auto UUID | Unique instance ID for this frontend | -| `--peer-address` | `LOCALAI_PEER_ADDRESS` | *(found from the route to the database)* | `host:port` at which the other frontends dial this one. A frontend publishes it so that a replica that does not hold the connection of a worker can reach the replica that does. Set it when the address found from the route to the database is not the one the other frontends should use. | +| `--peer-address` | `LOCALAI_PEER_ADDRESS` | *(found from the route to the database)* | `host:port` at which the other frontends dial this one. A frontend publishes it so that a replica that does not hold the connection of a worker can reach the replica that does. Set it when the address found from the route to the database is not the one the other frontends should use. The link carries the credential of the replica and every request that it relays, so by default it is clear text. Use `--peer-tls` when the replicas sit behind TLS. | +| `--peer-tls` | `LOCALAI_PEER_TLS` | `false` | Dial the other frontends over `wss` and not `ws`. Every frontend must then be reachable over TLS at the address it publishes with `--peer-address`, for example behind a reverse proxy that terminates TLS. LocalAI warns at start-up when it publishes an address that is not on the local host and this flag is off. | +| `--peer-tls-ca` | `LOCALAI_PEER_TLS_CA` | *(system roots)* | PEM file with the certificate authority that signs the certificate of the other frontends. Use it with a private CA. | | `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS server URL (e.g., `nats://localhost:4222`) | | `--registration-token` | `LOCALAI_REGISTRATION_TOKEN` | *(empty)* | Token that workers must provide to register | | `--registration-require-auth` | `LOCALAI_REGISTRATION_REQUIRE_AUTH` | `false` | Fail startup when distributed mode is enabled but the registration token is empty (node endpoints and worker file-transfer would otherwise be unauthenticated) | From 9871365662af4071173fc36abf5a65dabdfe26ea Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 07:57:42 +0000 Subject: [PATCH 096/162] fix(nodes): treat a stream that broke during an install as an install that may still run A streamed install, upgrade or forced install whose body ended or whose connection broke after the worker had answered 200 was a plain error. The caller counted it as a failed attempt and fired the install again while the first one could still be running. The NATS carrier already reads a wait that ran out as ErrWorkerStillInstalling. The tunnel carrier now does the same for a stream that broke after the worker accepted the verb. A reply line that is not JSON is still an ordinary error: the worker answered, out of protocol. The spec that held the old behaviour is rewritten for the new one. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/control_client.go | 13 +++++++++ core/services/nodes/control_tunnel.go | 11 +++++--- core/services/nodes/control_tunnel_test.go | 32 ++++++++++++++++++++-- 3 files changed, 49 insertions(+), 7 deletions(-) diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go index 7f90b2bf620f..38787ae4dc4d 100644 --- a/core/services/nodes/control_client.go +++ b/core/services/nodes/control_client.go @@ -16,6 +16,10 @@ import ( "github.com/mudler/LocalAI/pkg/httpclient" ) +// errStreamBroken marks the failure of a streamed verb after the worker accepted +// it: the body ended or the connection broke before the reply line. +var errStreamBroken = errors.New("the stream of the verb broke after the worker accepted it") + // ControlClient sends the control verbs of a frontend to a worker as HTTP // requests, over whatever dialer reaches that worker. In a tunnel deployment the // dialer opens a stream with the http tag, so the request reaches the HTTP @@ -172,6 +176,15 @@ func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, verb string, if errors.Is(err, io.EOF) { err = fmt.Errorf("the %s stream ended before its reply line: %w", verb, io.ErrUnexpectedEOF) } + // A body that ends, or a connection that breaks, after the worker + // answered 200 is a verb that was accepted and whose reply was lost. + // A line that is not JSON is the worker speaking out of protocol, which + // is another thing. + var syntax *json.SyntaxError + var typed *json.UnmarshalTypeError + if !errors.As(err, &syntax) && !errors.As(err, &typed) { + err = fmt.Errorf("%w: %w", errStreamBroken, err) + } return controlFailure(ctx, nodeID, err) } if env.Reply != nil { diff --git a/core/services/nodes/control_tunnel.go b/core/services/nodes/control_tunnel.go index eb3a1f4b932c..a217757fb6b3 100644 --- a/core/services/nodes/control_tunnel.go +++ b/core/services/nodes/control_tunnel.go @@ -84,10 +84,13 @@ func (l *httpLink) requestWithProgress(ctx context.Context, nodeID, verb, opID s return l.client.CallStreaming(ctx, nodeID, verb, req, reply, sink) } -// waitExpired is true for the timeout that this link set on the request. Nothing -// else reads as a wait that ran out: a stream that broke is not a worker that is -// still installing. -func (l *httpLink) waitExpired(err error) bool { return errors.Is(err, context.DeadlineExceeded) } +// waitExpired is true for the timeout that this link set on the request, and for +// a stream that broke after the worker accepted the verb. In both the work may +// still be running, and the caller must not count the call as a failed attempt +// and fire it again. A reply that cannot be read is neither. +func (l *httpLink) waitExpired(err error) bool { + return errors.Is(err, context.DeadlineExceeded) || errors.Is(err, errStreamBroken) +} // acknowledgementMissing is always false. An HTTP request has no way to tell a // worker that stopped the backend and did not reply from one that did not diff --git a/core/services/nodes/control_tunnel_test.go b/core/services/nodes/control_tunnel_test.go index 4fa19cf88793..2fe809747473 100644 --- a/core/services/nodes/control_tunnel_test.go +++ b/core/services/nodes/control_tunnel_test.go @@ -356,17 +356,43 @@ var _ = Describe("The tunnel carrier of the control verbs", func() { Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) }) - It("does not say that the worker may still be installing when the stream ends without a reply", func() { + It("says the worker may still be installing when the stream ends without a reply", func() { + // The worker accepted the install and the stream broke in the middle of + // it. That is the same situation as a wait that ran out on NATS: the + // work may be running. Counting it as a failed attempt would fire the + // install again. worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"progress":{"op_id":"x"}}`+"\n") }) _, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", nil) Expect(err).To(HaveOccurred()) - Expect(errors.Is(err, io.ErrUnexpectedEOF)).To(BeTrue()) - Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeFalse(), "the link broke; nothing was learned about the install") + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(), "%v", err) + Expect(err.Error()).To(ContainSubstring("before its reply line"), "the cause stays in the message") Expect(errors.Is(err, ErrNoRoute)).To(BeFalse()) }) + It("says the same of an upgrade and of the legacy install with Force", func() { + broken := func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"progress":{"op_id":"x"}}`+"\n") + } + worker.on(workerctl.VerbBackendUpgrade, broken) + _, err := control.UpgradeBackend(node, "b", "", "", "", "", 0, "op", nil) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(), "%v", err) + + worker.on(workerctl.VerbBackendInstall, broken) + _, err = control.InstallBackendForce(node, "b", "", "", "", "", 0, "op", nil) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(), "%v", err) + }) + + It("does not say that the worker may still be installing when its reply line is garbage", func() { + worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "this is not a line of the stream\n") + }) + _, err := control.InstallBackend(node, "b", "m", "", "", "", "", 0, "", nil) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeFalse(), "the worker spoke, and not in the protocol: %v", err) + }) + It("hands a refusal of the worker back in the reply", func() { worker.on(workerctl.VerbBackendInstall, func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, `{"reply":{"success":false,"error":"no space left"}}`+"\n") From 9b140d4478a3542b64d834a03ef2192bd1eda27e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:00:49 +0000 Subject: [PATCH 097/162] fix(nodes): probe a replica again before loading its model a second time A probe that could not reach the worker sent the warm path to a cold load at once. After a relay that dropped for a moment, this put a second replica of the model beside the one that was well. The warm path now repeats the probe once after a short wait, and falls through to a load elsewhere only when the transport is still down. The row of an unreachable replica stays. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/router.go | 18 ++++- core/services/nodes/transport_failure_test.go | 78 +++++++++++++++++++ 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index e28b679cdce6..1be2ef284f96 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -760,6 +760,11 @@ type routeAttempt struct { observeChain []uint64 } +// warmProbeRetryDelay is how long the warm path waits before it probes a replica +// again, when the first probe could not reach the worker. A variable, so that a +// spec does not wait for it. +var warmProbeRetryDelay = 500 * time.Millisecond + // tryWarmPath returns a route to an already-loaded, reachable replica, or nil // when the model has to be cold-loaded. It is the authority on readiness: a // waiter woken by a finished job re-runs it rather than trusting the signal, @@ -776,7 +781,18 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route replicaIdx := nm.ReplicaIndex // Verify the backend process is still alive via gRPC health check - switch r.probeHealth(ctx, node, modelAddr) { + verdict := r.probeHealth(ctx, node, modelAddr) + if verdict == probeUnknown { + // A relay that dropped for a moment looks the same as a worker that cannot + // be reached. Loading the model again at once would put a second replica + // beside one that is well, so the probe is repeated once first. + select { + case <-time.After(warmProbeRetryDelay): + verdict = r.probeHealth(ctx, node, modelAddr) + case <-ctx.Done(): + } + } + switch verdict { case probeUnknown: // The transport failed, so nothing is known about the backend. The row // stays: a replica that is loaded on a worker that cannot be reached from diff --git a/core/services/nodes/transport_failure_test.go b/core/services/nodes/transport_failure_test.go index 1ff0b913aef9..cbef409c6369 100644 --- a/core/services/nodes/transport_failure_test.go +++ b/core/services/nodes/transport_failure_test.go @@ -276,3 +276,81 @@ var _ = Describe("The router and a transport that is slow", func() { Expect(reg.removeCalls).ToNot(ContainElement("n-old:m"), "a transport that is slow is not a replica that is gone") }) }) + +// blipProbe is a backend whose first health checks fail in the transport and +// whose later ones succeed: a relay that dropped for a moment. +type blipProbe struct { + *stubBackend + mu sync.Mutex + blips int + calls int +} + +func (b *blipProbe) HealthCheck(context.Context) (bool, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.calls++ + if b.calls <= b.blips { + return false, errors.New("unavailable") + } + return true, nil +} + +func (b *blipProbe) LastDialError() error { + b.mu.Lock() + defer b.mu.Unlock() + if b.calls <= b.blips { + return errNoTunnel + } + return nil +} + +var _ = Describe("The warm path and a transport that drops for a moment", func() { + route := func(blips int) (*fakeModelRouter, *fakeUnloader, *blipProbe) { + GinkgoHelper() + DeferCleanup(func(old time.Duration) { warmProbeRetryDelay = old }, warmProbeRetryDelay) + warmProbeRetryDelay = 10 * time.Millisecond + probe := &blipProbe{stubBackend: &stubBackend{loadResult: &pb.Result{Success: true}}, blips: blips} + reg := &fakeModelRouter{ + findAndLockNode: &BackendNode{ID: "n-old", Name: "old-node", Address: "10.0.0.70:50051"}, + findAndLockNM: &NodeModel{NodeID: "n-old", ModelName: "m", Address: "10.0.0.70:9001"}, + findIdleNode: &BackendNode{ID: "n-new", Name: "new-node", Address: "10.0.0.71:50051"}, + } + unloader := &fakeUnloader{installReply: &workerctl.BackendInstallReply{Success: true, Address: "10.0.0.71:9001"}} + router := NewSmartRouter(&warmOnce{fakeModelRouter: reg}, SmartRouterOptions{Unloader: unloader, ClientFactory: &sameClient{probe}}) + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) + Expect(err).ToNot(HaveOccurred()) + return reg, unloader, probe + } + + It("probes again before it loads the model a second time", func() { + reg, unloader, _ := route(1) + Expect(unloader.installCalls).To(BeEmpty(), "the replica was well; a second load would duplicate it") + Expect(reg.removeCalls).To(BeEmpty()) + }) + + It("still falls through to a load elsewhere when the transport stays down", func() { + reg, unloader, _ := route(10) + Expect(unloader.installCalls).ToNot(BeEmpty()) + Expect(reg.removeCalls).ToNot(ContainElement("n-old:m"), "the row of an unreachable replica stays") + }) +}) + +type sameClient struct{ c grpc.Backend } + +func (s *sameClient) NewClient(_, _ string, _ bool) grpc.Backend { return s.c } + +// warmOnce finds the loaded replica for the first lookup only, so that a probe +// that gives up leads to a load and not to a second pass over the warm path. +type warmOnce struct { + *fakeModelRouter + asked bool +} + +func (w *warmOnce) FindAndLockNodeWithModel(ctx context.Context, name string, ids []string, pref *RoutePreference) (*BackendNode, *NodeModel, error) { + if w.asked { + return nil, nil, nil + } + w.asked = true + return w.fakeModelRouter.FindAndLockNodeWithModel(ctx, name, ids, pref) +} From 1e9fd39e5be283b06920a893e2ce0307a4a0024b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:00:49 +0000 Subject: [PATCH 098/162] fix(worker): keep a control verb running when its caller goes away The HTTP control server passed the context of the request to the handlers, and the context ends when the caller hangs up. The NATS server passes a context that nothing cancels. An unload whose frontend gave up on the reply stopped half way through freeing a backend. Both servers now behave the same way: the handlers get a context that the caller cannot cancel and bound their own work. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/control_http.go | 10 ++- .../worker/control_http_context_test.go | 78 +++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) create mode 100644 core/services/worker/control_http_context_test.go diff --git a/core/services/worker/control_http.go b/core/services/worker/control_http.go index 3829481371b9..0efddcb064e0 100644 --- a/core/services/worker/control_http.go +++ b/core/services/worker/control_http.go @@ -1,6 +1,7 @@ package worker import ( + "context" "encoding/json" "fmt" "net/http" @@ -63,6 +64,11 @@ func (s *httpControlServer) route(v controlVerb) (string, error) { return workerctl.PathOf(string(v)), nil } +// The handlers of a verb get a context that the caller going away does not +// cancel, as the NATS server does with context.Background. A verb bounds its own +// work, and an unload that stopped half way because the frontend gave up on the +// reply would leave a backend half freed. + // handle serves a verb that answers with one JSON body, or with 204 when the // handler returns no reply. Requests of one verb are delivered concurrently: // the NATS server delivers them one at a time, and the handlers are written to @@ -77,7 +83,7 @@ func (s *httpControlServer) handle(v controlVerb, h controlHandler) error { if !ok { return } - reply, undecodable := h(r.Context(), body) + reply, undecodable := h(context.WithoutCancel(r.Context()), body) if undecodable != nil { // The body is not a request of this verb. The typed refusal in reply // is what the NATS server sends. Here the request failed before it @@ -112,7 +118,7 @@ func (s *httpControlServer) handleWithProgress(v controlVerb, h progressControlH return } stream := newNDJSONStream(w) - reply, undecodable := h(r.Context(), body, stream.progress) + reply, undecodable := h(context.WithoutCancel(r.Context()), body, stream.progress) if undecodable != nil { // Nothing was written yet unless a handler reported progress before it // decoded, which none does. The status is still open. diff --git a/core/services/worker/control_http_context_test.go b/core/services/worker/control_http_context_test.go new file mode 100644 index 000000000000..47caa8eb76fb --- /dev/null +++ b/core/services/worker/control_http_context_test.go @@ -0,0 +1,78 @@ +package worker + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// The NATS server hands a handler a context that nothing cancels. A verb such as +// an unload bounds its own work, and it must run to the end when the caller goes +// away, so that the HTTP server behaves in the same way. +var _ = Describe("The context of a verb over HTTP", func() { + // callAndLeave sends a request and cancels it as soon as the handler runs. + callAndLeave := func(srv *httptest.Server, path string, started <-chan struct{}) { + GinkgoHelper() + reqCtx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, srv.URL+path, strings.NewReader("{}")) + Expect(err).ToNot(HaveOccurred()) + go func() { + resp, err := http.DefaultClient.Do(req) + if err == nil { + _ = resp.Body.Close() + } + }() + Eventually(started).Should(BeClosed()) + cancel() + // Long enough for the server to see the connection close. + time.Sleep(200 * time.Millisecond) + } + + It("is not cancelled when the caller goes away, for a verb with one reply", func() { + ctl := newHTTPControlServer() + started, finish := make(chan struct{}), make(chan struct{}) + seen := make(chan error, 1) + Expect(ctl.handle(verbModelUnload, func(ctx context.Context, _ []byte) (any, error) { + close(started) + <-finish + seen <- ctx.Err() + return workerctl.ModelUnloadReply{Success: true}, nil + })).To(Succeed()) + srv := httptest.NewServer(ctl) + DeferCleanup(srv.Close) + + callAndLeave(srv, workerctl.PathOf(workerctl.VerbModelUnload), started) + close(finish) + var got error + Eventually(seen).Should(Receive(&got)) + Expect(got).ToNot(HaveOccurred(), "the free of a backend must not be abandoned half way") + }) + + It("is not cancelled when the caller goes away, for a verb that reports progress", func() { + ctl := newHTTPControlServer() + started, finish := make(chan struct{}), make(chan struct{}) + seen := make(chan error, 1) + Expect(ctl.handleWithProgress(verbBackendInstall, func(ctx context.Context, _ []byte, _ progressSink) (any, error) { + close(started) + <-finish + seen <- ctx.Err() + return workerctl.BackendInstallReply{Success: true}, nil + })).To(Succeed()) + srv := httptest.NewServer(ctl) + DeferCleanup(srv.Close) + + callAndLeave(srv, workerctl.PathOf(workerctl.VerbBackendInstall), started) + close(finish) + var got error + Eventually(seen).Should(Receive(&got)) + Expect(got).ToNot(HaveOccurred()) + }) +}) From 27bc1160104a8b6fd70daeea5690d91a0f61eb15 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:00:49 +0000 Subject: [PATCH 099/162] fix(carrier): give back the idle streams of the tunnel carrier when its set closes The control client and the HTTP file stager of the tunnel set keep one HTTP client for each node, with idle streams on the tunnel. Closing the set closed only the fan-out, so the streams stayed open until their idle timeout. Both clients now have a Close that drops every client and closes the idle streams, and the set calls them. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/tunnel.go | 15 +++- core/services/nodes/client_close_test.go | 91 ++++++++++++++++++++++++ core/services/nodes/control_client.go | 17 +++++ core/services/nodes/file_stager_http.go | 15 ++++ 4 files changed, 136 insertions(+), 2 deletions(-) create mode 100644 core/services/nodes/client_close_test.go diff --git a/core/services/carrier/tunnel.go b/core/services/carrier/tunnel.go index c122c1c70709..ea3b686a076d 100644 --- a/core/services/carrier/tunnel.go +++ b/core/services/carrier/tunnel.go @@ -79,6 +79,7 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { var ( files FileCarrier forgetFile func(string) + closeFiles func() ) if o.S3Staging { if o.FileManager == nil { @@ -93,7 +94,7 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { // the node. It names the node in a log line and in the URL. return nodes.WorkerHTTPHost(nodeID, ""), nil }, o.Token, bulkDial) - files, forgetFile = stager, stager.ForgetNode + files, forgetFile, closeFiles = stager, stager.ForgetNode, stager.Close } set := &Set{ @@ -114,7 +115,17 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { forgetFile(nodeID) } }, - Close: o.Fanout.Close, + // The streams that the clients keep idle belong to tunnels that are about to + // be replaced, so they are given back with the set. + Close: func() { + control.Close() + if closeFiles != nil { + closeFiles() + } + if o.Fanout.Close != nil { + o.Fanout.Close() + } + }, } if err := set.Validate(); err != nil { return nil, err diff --git a/core/services/nodes/client_close_test.go b/core/services/nodes/client_close_test.go new file mode 100644 index 000000000000..c12c90c60ece --- /dev/null +++ b/core/services/nodes/client_close_test.go @@ -0,0 +1,91 @@ +package nodes + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// The clients of a tunnel carrier keep idle streams of the tunnel. When the +// carrier is closed they must give them back, and not wait for the idle timeout +// of 90 seconds. +var _ = Describe("Closing the clients of the tunnel carrier", func() { + // server counts the connections it has open. + server := func() (*httptest.Server, *atomic.Int32) { + GinkgoHelper() + var open atomic.Int32 + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte("{}")) + })) + srv.Config.ConnState = func(_ net.Conn, s http.ConnState) { + switch s { + case http.StateNew: + open.Add(1) + case http.StateClosed, http.StateHijacked: + open.Add(-1) + } + } + srv.Start() + DeferCleanup(srv.Close) + return srv, &open + } + + dialTo := func(srv *httptest.Server) WorkerNetDialerFor { + return func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", srv.Listener.Addr().String()) + } + } + } + + It("drops the idle streams of every node held by the control client", func() { + srv, open := server() + client := NewControlClient(dialTo(srv), "") + for _, node := range []string{"n1", "n2"} { + Expect(client.Call(context.Background(), node, workerctl.VerbBackendList, struct{}{}, &workerctl.BackendListReply{})).To(Succeed()) + } + Eventually(open.Load).Should(Equal(int32(2)), "one idle stream for each node") + + client.Close() + Eventually(open.Load, 5*time.Second).Should(BeZero()) + }) + + It("drops the idle streams of every node held by the file stager", func() { + srv, open := server() + stager := NewHTTPFileStager(func(string) (string, error) { return srv.Listener.Addr().String(), nil }, "", dialTo(srv)) + for _, node := range []string{"n1", "n2"} { + resp, err := mustClient(stager, node).Get(srv.URL) + Expect(err).ToNot(HaveOccurred()) + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + Eventually(open.Load).Should(Equal(int32(2))) + + stager.Close() + Eventually(open.Load, 5*time.Second).Should(BeZero()) + }) + + It("is safe on nil and when called twice", func() { + var c *ControlClient + var s *HTTPFileStager + Expect(func() { c.Close(); s.Close() }).ToNot(Panic()) + client := NewControlClient(nil, "") + Expect(func() { client.Close(); client.Close() }).ToNot(Panic()) + }) +}) + +func mustClient(s *HTTPFileStager, node string) *http.Client { + c, err := s.clientFor(node) + Expect(err).ToNot(HaveOccurred()) + return c +} diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go index 38787ae4dc4d..7ec3e6b9c68a 100644 --- a/core/services/nodes/control_client.go +++ b/core/services/nodes/control_client.go @@ -126,6 +126,23 @@ func (c *ControlClient) ForgetNode(nodeID string) { } } +// Close drops the client of every node and closes the idle streams that their +// transports hold. The carrier calls it when it is replaced, so that the streams +// of its tunnels do not wait for an idle timeout. It is safe on a nil receiver +// and may be called twice. +func (c *ControlClient) Close() { + if c == nil { + return + } + c.clientsMu.Lock() + entries := c.clients + c.clients = map[string]*nodeHTTPClient{} + c.clientsMu.Unlock() + for _, entry := range entries { + entry.transport.CloseIdleConnections() + } +} + // Call sends one request and decodes the reply. reply may be nil for a verb that // answers with no body. // diff --git a/core/services/nodes/file_stager_http.go b/core/services/nodes/file_stager_http.go index 23c681ef5849..e56c03950f0a 100644 --- a/core/services/nodes/file_stager_http.go +++ b/core/services/nodes/file_stager_http.go @@ -142,6 +142,21 @@ func (h *HTTPFileStager) ForgetNode(nodeID string) { } } +// Close drops the client of every node and closes the idle streams that their +// transports hold. It is safe on a nil receiver and may be called twice. +func (h *HTTPFileStager) Close() { + if h == nil { + return + } + h.clientsMu.Lock() + entries := h.clients + h.clients = map[string]*nodeFileClient{} + h.clientsMu.Unlock() + for _, entry := range entries { + entry.transport.CloseIdleConnections() + } +} + // ReleaseRemote removes one exact ephemeral key from a backend node. func (h *HTTPFileStager) ReleaseRemote(ctx context.Context, nodeID, key string) error { if err := validateEphemeralReleaseKey(key); err != nil { From f2127402266710dc395d1fcd860b92b59777e90e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:04:26 +0000 Subject: [PATCH 100/162] refactor(nodes): remove StopNode, which nothing calls RemoteUnloaderAdapter.StopNode published the node stop subject, and no production code called it. The worker still serves the node.stop verb on both carriers. The spec for the method and the end-to-end context that only exercised it are gone too. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/unloader.go | 6 ----- core/services/nodes/unloader_test.go | 9 -------- tests/e2e/distributed/node_lifecycle_test.go | 23 -------------------- 3 files changed, 38 deletions(-) diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index cab5160afca8..cbfe6a8873e7 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -655,12 +655,6 @@ func (a *nodeControl) DeleteModelFiles(modelName string) error { return nil } -// StopNode tells a worker node to shut down entirely (deregister + exit). -func (a *RemoteUnloaderAdapter) StopNode(nodeID string) error { - subject := messaging.SubjectNodeStop(nodeID) - return a.nats.Publish(subject, nil) -} - // errVerbNotServed means that the worker answered and does not serve the verb. // It is a fact about the version of the worker. It is not a routing fact and // not an answer about any backend, so only a caller that has a fallback for an diff --git a/core/services/nodes/unloader_test.go b/core/services/nodes/unloader_test.go index d3d13e269995..8dbefc6aaddd 100644 --- a/core/services/nodes/unloader_test.go +++ b/core/services/nodes/unloader_test.go @@ -363,15 +363,6 @@ var _ = Describe("RemoteUnloaderAdapter", func() { }) }) - Describe("StopNode", func() { - It("publishes to correct subject", func() { - Expect(adapter.StopNode("node-abc")).To(Succeed()) - Expect(mc.published).To(HaveLen(1)) - Expect(mc.published[0].Subject).To(Equal(messaging.SubjectNodeStop("node-abc"))) - Expect(mc.published[0].Data).To(BeNil()) - }) - }) - Describe("DeleteModelFiles", func() { It("with no nodes returns nil", func() { locator.nodes = nil diff --git a/tests/e2e/distributed/node_lifecycle_test.go b/tests/e2e/distributed/node_lifecycle_test.go index 2ca7707c1c4c..54067d4eac5c 100644 --- a/tests/e2e/distributed/node_lifecycle_test.go +++ b/tests/e2e/distributed/node_lifecycle_test.go @@ -156,29 +156,6 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f }) }) - Context("NATS node stop events (full shutdown)", func() { - It("should publish stop event to a node", func() { - node := &nodes.BackendNode{ - Name: "stop-me", Address: "h3:50051", - } - Expect(registry.Register(context.Background(), node, true)).To(Succeed()) - - var stopped atomic.Int32 - sub, err := infra.NC.Subscribe(messaging.SubjectNodeStop(node.ID), func(data []byte) { - stopped.Add(1) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - FlushNATS(infra.NC) - - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) - Expect(adapter.StopNode(node.ID)).To(Succeed()) - - Eventually(func() int32 { return stopped.Load() }, "5s").Should(Equal(int32(1))) - }) - }) - Context("NATS subject naming", func() { It("should generate correct backend lifecycle subjects", func() { Expect(messaging.SubjectNodeBackendInstall("node-abc")).To(Equal("nodes.node-abc.backend.install")) From 0dcd860f201ce9fbd7193b33e4d9fb1f21d83573 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:04:26 +0000 Subject: [PATCH 101/162] test(worker): split the process kill helper so that the tests build on Windows operations_test.go called syscall.Kill, which does not exist on Windows, so go vet failed for the test build of the package. The helper now lives in one file for Unix and one for Windows. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/kill_unix_test.go | 8 ++++++++ core/services/worker/kill_windows_test.go | 14 ++++++++++++++ core/services/worker/operations_test.go | 3 --- 3 files changed, 22 insertions(+), 3 deletions(-) create mode 100644 core/services/worker/kill_unix_test.go create mode 100644 core/services/worker/kill_windows_test.go diff --git a/core/services/worker/kill_unix_test.go b/core/services/worker/kill_unix_test.go new file mode 100644 index 000000000000..abb27aebbf37 --- /dev/null +++ b/core/services/worker/kill_unix_test.go @@ -0,0 +1,8 @@ +//go:build !windows + +package worker + +import "syscall" + +// syscallKill ends a process the way a crash would, with no chance to clean up. +func syscallKill(pid int) error { return syscall.Kill(pid, syscall.SIGKILL) } diff --git a/core/services/worker/kill_windows_test.go b/core/services/worker/kill_windows_test.go new file mode 100644 index 000000000000..4ed2eef0fb7f --- /dev/null +++ b/core/services/worker/kill_windows_test.go @@ -0,0 +1,14 @@ +//go:build windows + +package worker + +import "os" + +// syscallKill ends a process the way a crash would, with no chance to clean up. +func syscallKill(pid int) error { + p, err := os.FindProcess(pid) + if err != nil { + return err + } + return p.Kill() +} diff --git a/core/services/worker/operations_test.go b/core/services/worker/operations_test.go index afaf1352cf13..c887af753009 100644 --- a/core/services/worker/operations_test.go +++ b/core/services/worker/operations_test.go @@ -8,7 +8,6 @@ import ( "strconv" "strings" "sync/atomic" - "syscall" "time" process "github.com/mudler/go-processmanager" @@ -328,8 +327,6 @@ var _ = Describe("model.unload", func() { func grandAlive(pid int) bool { return pidAlive(strconv.Itoa(pid)) } -func syscallKill(pid int) error { return syscall.Kill(pid, syscall.SIGKILL) } - func startRegisteredBackend(register func(*gogrpc.Server)) (string, int, func()) { lis, err := net.Listen("tcp", "127.0.0.1:0") Expect(err).NotTo(HaveOccurred()) From 411d5735c6a996f3c36fd7a3433fa28ba3687a9c Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:04:26 +0000 Subject: [PATCH 102/162] test(carrier): judge the probe on a busy lane against the idle lane The spec that keeps a small call fast during a transfer compared a 99th percentile of at least ten probes with a fixed 500 ms. On a loaded machine one slow probe decided the result. It now measures the idle lane first, uses the 95th percentile of the probes, needs more probes, and bounds the busy lane by a multiple of the idle lane. The comparison with the transfer on the shared lane stays. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../carrier/tunnel_integration_test.go | 59 ++++++++++++++----- 1 file changed, 44 insertions(+), 15 deletions(-) diff --git a/core/services/carrier/tunnel_integration_test.go b/core/services/carrier/tunnel_integration_test.go index 7d9ec3112a8f..4b40d279a721 100644 --- a/core/services/carrier/tunnel_integration_test.go +++ b/core/services/carrier/tunnel_integration_test.go @@ -526,10 +526,38 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { time.Sleep(time.Second) src := filepath.Join(GinkgoT().TempDir(), "big.bin") - Expect(os.WriteFile(src, make([]byte, 16<<20), 0o600)).To(Succeed()) + Expect(os.WriteFile(src, make([]byte, 32<<20), 0o600)).To(Succeed()) + + // percentile is the 95th percentile of a set of samples. The 99th of a few + // dozen samples is the largest one, which a single scheduling hiccup of a + // loaded machine decides. + percentile := func(samples []time.Duration) time.Duration { + sorted := slices.Clone(samples) + slices.Sort(sorted) + return sorted[min(len(sorted)-1, len(sorted)*95/100)] + } + + // probe is one call on the inference lane. + probe := func() (time.Duration, error) { + begun := time.Now() + _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + return time.Since(begun), err + } + + // The baseline of the lane: the same probe while nothing else runs. Both + // runs below are judged against it, so that a slow machine moves the bound + // with it. + var idle []time.Duration + for range 60 { + took, err := probe() + Expect(err).ToNot(HaveOccurred()) + idle = append(idle, took) + time.Sleep(10 * time.Millisecond) + } + idleP95 := percentile(idle) // probeDuring probes the inference lane while transfer runs, and returns - // the 99th percentile of the probes and the time of the transfer. + // the 95th percentile of the probes and the time of the transfer. probeDuring := func(transfer func() error) (time.Duration, time.Duration) { GinkgoHelper() var ( @@ -544,11 +572,11 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { return case <-time.After(10 * time.Millisecond): } - begun := time.Now() - if _, err := a.commands.OperationControl(nodeID, workerctl.OperationRequest{Renew: []string{"op"}}); err != nil { + took, err := probe() + if err != nil { return } - samples = append(samples, time.Since(begun)) + samples = append(samples, took) } }) begun := time.Now() @@ -558,13 +586,12 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { wg.Wait() Expect(err).ToNot(HaveOccurred()) Expect(elapsed).To(BeNumerically(">", 500*time.Millisecond), "the transfer must last long enough to be measured") - Expect(len(samples)).To(BeNumerically(">=", 10), "probes completed during the transfer") - slices.Sort(samples) - return samples[min(len(samples)-1, len(samples)*99/100)], elapsed + Expect(len(samples)).To(BeNumerically(">=", 30), "probes completed during the transfer: %d in %v", len(samples), elapsed) + return percentile(samples), elapsed } // The transfer on the bulk lane, as the set does it. - bulkP99, bulkTime := probeDuring(func() error { + bulkP95, bulkTime := probeDuring(func() error { _, err := a.files.EnsureRemote(ctx, nodeID, src, "models/big.bin") return err }) @@ -572,17 +599,19 @@ var _ = Describe("The tunnel carrier through the holders, end to end", func() { // would do if it fell back. The comparison is between two runs on the // same machine, so a loaded runner moves both. shared := nodes.NewHTTPFileStager(func(id string) (string, error) { return nodes.WorkerHTTPHost(id, ""), nil }, registrationToken, a.set.Dialer) - sharedP99, sharedTime := probeDuring(func() error { + sharedP95, sharedTime := probeDuring(func() error { _, err := shared.EnsureRemote(ctx, nodeID, src, "models/big-shared.bin") return err }) AddReportEntry("probe during a transfer", map[string]any{ - "bulk lane p99 ms": bulkP99.Milliseconds(), "bulk transfer s": bulkTime.Seconds(), - "shared lane p99 ms": sharedP99.Milliseconds(), "shared transfer s": sharedTime.Seconds(), + "idle lane p95 ms": idleP95.Milliseconds(), + "bulk lane p95 ms": bulkP95.Milliseconds(), "bulk transfer s": bulkTime.Seconds(), + "shared lane p95 ms": sharedP95.Milliseconds(), "shared transfer s": sharedTime.Seconds(), }) - Expect(bulkP99).To(BeNumerically("<", sharedP99*8/10), - "a probe must wait less behind a transfer on the bulk lane than behind one on its own lane: bulk %v, shared %v", bulkP99, sharedP99) - Expect(bulkP99).To(BeNumerically("<", 500*time.Millisecond), "a sanity bound for the probe on a loaded machine") + Expect(bulkP95).To(BeNumerically("<", sharedP95*8/10), + "a probe must wait less behind a transfer on the bulk lane than behind one on its own lane: bulk %v, shared %v", bulkP95, sharedP95) + Expect(bulkP95).To(BeNumerically("<", idleP95*5+100*time.Millisecond), + "a transfer on the bulk lane must not slow a probe much beyond the idle lane: idle %v, bulk %v", idleP95, bulkP95) }) It("refuses a transfer when the worker has no bulk session, and sends nothing on the other lane", func() { From d667b1b41ae425e3a73ebb53bb0b3c2a3742c130 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 08:10:02 +0000 Subject: [PATCH 103/162] feat(worker): let a worker on NATS attach to a tunnel while it runs A worker that booted on NATS never mounted the HTTP control plane and never started a tunnel, so a later change of the cluster to the tunnel needed a restart of every worker. The control plane is now a slot on the HTTP server of the worker (nodes.ControlSlot). While it is empty the path answers 404, as before, so a worker on NATS behaves as it did. tunnelFollower fills the slot, claims the worker verbs on it once, and starts the tunnel; Detach reverses it, and a worker that has no frontend URL to dial does not mount a slot at all. The bearer check, or the loopback rule when there is no token, applies to the attached plane. A worker that boots on the tunnel starts it as before. Nothing calls the follower yet: the decision to follow belongs to a later change. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/file_transfer_server.go | 50 +++++- core/services/worker/follow.go | 152 ++++++++++++++++ core/services/worker/follow_test.go | 185 ++++++++++++++++++++ core/services/worker/supervisor.go | 5 + core/services/worker/worker.go | 67 ++++--- 5 files changed, 436 insertions(+), 23 deletions(-) create mode 100644 core/services/worker/follow.go create mode 100644 core/services/worker/follow_test.go diff --git a/core/services/nodes/file_transfer_server.go b/core/services/nodes/file_transfer_server.go index bce11ee4d447..9f5fd1b947b6 100644 --- a/core/services/nodes/file_transfer_server.go +++ b/core/services/nodes/file_transfer_server.go @@ -17,6 +17,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/gorilla/websocket" @@ -182,7 +183,11 @@ func startFileTransferServerWithControl(lis net.Listener, stagingDir, modelsDir, // a tunnel. The bearer check is the one of the file routes: a verb can stop a // node, so it is never served without it. if control != nil { - mux.Handle(workerctl.Prefix, controlGate(token, control)) + if slot, ok := control.(*ControlSlot); ok { + mux.Handle(workerctl.Prefix, slot.gated(token)) + } else { + mux.Handle(workerctl.Prefix, controlGate(token, control)) + } } // Backend log endpoints (only registered when a log store is provided) @@ -1080,6 +1085,49 @@ func resolveKeyToDir(key, stagingDir, modelsDir, dataDir string) (targetDir, rel return } +// ControlSlot is a control plane that a worker can attach after its HTTP server +// started, and detach again. A worker that booted on NATS has none, and a later +// switch to the tunnel needs one without a restart. +// +// While nothing is attached, the path answers as a path that does not exist, so a +// worker that never attaches behaves as it did before the slot existed. The +// bearer check applies to what is attached, as for any control plane. +type ControlSlot struct { + handler atomic.Pointer[http.Handler] +} + +// Set attaches the control plane. A nil handler detaches it. +func (s *ControlSlot) Set(h http.Handler) { + if h == nil { + s.handler.Store(nil) + return + } + s.handler.Store(&h) +} + +// ServeHTTP serves the attached control plane, and 404 when there is none. +func (s *ControlSlot) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if h := s.handler.Load(); h != nil { + (*h).ServeHTTP(w, r) + return + } + http.NotFound(w, r) +} + +// gated is the slot behind controlGate. The check of the token comes after the +// test for an attached plane, so that a path with nothing behind it does not +// answer 401 to a caller who could not know that it exists. +func (s *ControlSlot) gated(token string) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + h := s.handler.Load() + if h == nil { + http.NotFound(w, r) + return + } + controlGate(token, *h).ServeHTTP(w, r) + }) +} + // controlGate puts the check of the control plane in front of next. // // The file routes have always fail open on an empty token, and a worker on NATS diff --git a/core/services/worker/follow.go b/core/services/worker/follow.go new file mode 100644 index 000000000000..f4eabd3f73fa --- /dev/null +++ b/core/services/worker/follow.go @@ -0,0 +1,152 @@ +package worker + +import ( + "context" + "errors" + "fmt" + "sync" + + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/xlog" +) + +// tunnelFollower attaches a worker to a tunnel, and detaches it, while the worker +// runs. +// +// A worker that booted on NATS has no tunnel and no control plane on its HTTP +// server. When the cluster moves to the tunnel the worker has to gain both +// without a restart, and this is the part that does it. What decides to attach is +// not here: the worker follows the carrier that the frontend names, and that +// logic is the caller's. +// +// The control plane is a slot on the HTTP server that the worker already runs. +// While the follower is detached the slot is empty and the path answers as a path +// that does not exist, which is what a worker on NATS has always answered. The +// verbs are claimed on the plane once, on the first attach: a plane can claim a +// verb only once, and a second attach must not fail on it. +type tunnelFollower struct { + cfg *Config + nodeID string + httpAddr string + token func() string + reauthorize func(ctx context.Context) error + slot *nodes.ControlSlot + register func(controlServer) error + control *httpControlServer + + mu sync.Mutex + tunnel *Tunnel + registered bool +} + +// tunnelFollowerOptions are the inputs of newTunnelFollower. +type tunnelFollowerOptions struct { + Config *Config + NodeID string + // HTTPAddr is where the HTTP server of this worker listens. The tunnel routes + // the http tag to it. + HTTPAddr string + // Token returns the own tunnel credential of the node, and Reauthorize mints + // a new one. See TunnelConfig. + Token func() string + Reauthorize func(ctx context.Context) error + // Slot is the control plane of the HTTP server. It is filled on attach and + // emptied on detach. + Slot *nodes.ControlSlot + // Register claims the verbs of the worker on a control plane. + Register func(controlServer) error +} + +func newTunnelFollower(o tunnelFollowerOptions) *tunnelFollower { + return &tunnelFollower{ + cfg: o.Config, + nodeID: o.NodeID, + httpAddr: o.HTTPAddr, + token: o.Token, + reauthorize: o.Reauthorize, + slot: o.Slot, + register: o.Register, + control: newHTTPControlServer(), + } +} + +// CanAttach reports whether this worker has what a tunnel needs: a frontend URL +// that can be turned into the endpoint of the tunnel. A worker that cannot attach +// must not mount a control plane at all. +func (f *tunnelFollower) CanAttach() bool { + if f == nil || f.cfg == nil { + return false + } + _, err := tunnelEndpoint(f.cfg.RegisterTo, f.nodeID) + return err == nil +} + +// Attached reports whether the tunnel is started. +func (f *tunnelFollower) Attached() bool { + f.mu.Lock() + defer f.mu.Unlock() + return f.tunnel != nil +} + +// Attach claims the verbs on the control plane if it has not, fills the slot, and +// starts the tunnel. It does nothing when the tunnel is started. The tunnel keeps +// trying to connect in the background, so a frontend that is down is not an +// error here. +func (f *tunnelFollower) Attach(ctx context.Context) error { + f.mu.Lock() + defer f.mu.Unlock() + if f.tunnel != nil { + return nil + } + if !f.CanAttach() { + return errors.New("attaching to the tunnel: this worker has no frontend URL to dial") + } + if !f.registered { + if err := f.register(f.control); err != nil { + return fmt.Errorf("serving the worker verbs over HTTP: %w", err) + } + f.registered = true + } + f.slot.Set(f.control) + t, err := StartTunnel(ctx, tunnelConfig(f.cfg, f.nodeID, f.httpAddr, f.token, f.reauthorize)) + if err != nil { + f.slot.Set(nil) + return fmt.Errorf("starting the worker tunnel: %w", err) + } + f.tunnel = t + xlog.Info("The worker attached to the tunnel", "node", f.nodeID) + return nil +} + +// Detach stops the tunnel and empties the slot. The verbs stay claimed on the +// plane for the next attach. It does nothing when the tunnel is not started. +func (f *tunnelFollower) Detach() error { + f.mu.Lock() + t := f.tunnel + f.tunnel = nil + f.mu.Unlock() + if t == nil { + return nil + } + f.slot.Set(nil) + err := t.Close() + xlog.Info("The worker detached from the tunnel", "node", f.nodeID) + return err +} + +// tunnelConfig is the configuration of the tunnel of a worker. A worker that +// boots on the tunnel and one that attaches later start it the same way. +func tunnelConfig(cfg *Config, nodeID, httpAddr string, token func() string, reauthorize func(ctx context.Context) error) TunnelConfig { + return TunnelConfig{ + FrontendURL: cfg.RegisterTo, + NodeID: nodeID, + Token: token, + // A frontend that refuses the credential is told who this worker is + // again, so that a rotated credential does not leave it unreachable. + Reauthorize: reauthorize, + // Built by tunnelServices and not inline, so that the routing table, + // which is the security boundary of the tunnel, can be reached from a + // spec without starting a worker. + Services: tunnelServices(cfg, httpAddr), + } +} diff --git a/core/services/worker/follow_test.go b/core/services/worker/follow_test.go new file mode 100644 index 000000000000..d5fbf07a4e33 --- /dev/null +++ b/core/services/worker/follow_test.go @@ -0,0 +1,185 @@ +package worker + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "strings" + "time" + + "github.com/labstack/echo/v4" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/workerctl" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// A worker that booted on NATS can be told, while it runs, to attach to a +// tunnel and serve its control verbs over it. These specs run the follower +// against the real connect endpoint, the real registry of tunnels and a real +// PostgreSQL. What decides to attach is not part of them. +var _ = Describe("A worker that attaches to a tunnel at run time", func() { + const ( + nodeToken = "own-tunnel-token" + registrationToken = "registration-token" + ) + + var ( + ctx context.Context + registry *tunnel.Registry + frontend *httptest.Server + nodeID string + httpAddr string + slot *nodes.ControlSlot + follower *tunnelFollower + ) + + freeAddr := func() string { + GinkgoHelper() + l, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + addr := l.Addr().String() + Expect(l.Close()).To(Succeed()) + return addr + } + + BeforeEach(func() { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(context.Background()) + DeferCleanup(cancel) + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR := cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + registry = tunnel.NewRegistry(clusterR, "replica-a") + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + nodeID = node.ID + sum := sha256.Sum256([]byte(nodeToken)) + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hex.EncodeToString(sum[:]))).To(Succeed()) + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(nodeReg, registry)) + frontend = httptest.NewServer(e) + DeferCleanup(frontend.Close) + + // The HTTP server of the worker, started as a worker on NATS starts it: the + // control plane is a slot with nothing in it. + httpAddr = freeAddr() + dir := GinkgoT().TempDir() + slot = &nodes.ControlSlot{} + srv, err := nodes.StartFileTransferServerWithControl(httpAddr, dir, dir, dir, registrationToken, 1<<20, nil, nil, slot) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { nodes.ShutdownFileTransferServer(srv) }) + Eventually(func() error { + c, err := net.DialTimeout("tcp", httpAddr, time.Second) + if err == nil { + _ = c.Close() + } + return err + }).Should(Succeed()) + + ss := newLifecycleTestSupervisor(make(chan os.Signal, 1)) + follower = newTunnelFollower(tunnelFollowerOptions{ + Config: &Config{RegisterTo: frontend.URL, ServeAddr: "127.0.0.1:50051"}, + NodeID: nodeID, + HTTPAddr: httpAddr, + Token: func() string { return nodeToken }, + Slot: slot, + Register: ss.registerLifecycleVerbs, + }) + DeferCleanup(func() { _ = follower.Detach() }) + }) + + direct := func() int { + GinkgoHelper() + req, err := http.NewRequest(http.MethodPost, "http://"+httpAddr+workerctl.PathOf(workerctl.VerbBackendList), strings.NewReader("{}")) + Expect(err).ToNot(HaveOccurred()) + req.Header.Set("Authorization", "Bearer "+registrationToken) + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + return resp.StatusCode + } + + // viaTunnel sends the control verb the way a frontend does: a stream with the + // http tag, and an HTTP request on it. + viaTunnel := func(token string) (int, string) { + GinkgoHelper() + dial := tunnel.NewWorkerDialer(registry, nil).DialerFor(nodeID, tunnel.StreamTagHTTP) + client := &http.Client{Transport: &http.Transport{DialContext: dial}, Timeout: 10 * time.Second} + DeferCleanup(client.CloseIdleConnections) + req, err := http.NewRequest(http.MethodPost, "http://"+nodeID+".worker.invalid"+workerctl.PathOf(workerctl.VerbBackendList), strings.NewReader("{}")) + Expect(err).ToNot(HaveOccurred()) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := client.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + body, _ := io.ReadAll(resp.Body) + return resp.StatusCode, string(body) + } + + It("serves no control path and holds no tunnel until it is told to attach", func() { + Expect(direct()).To(Equal(http.StatusNotFound), "a worker on NATS answers as it did before") + Expect(follower.Attached()).To(BeFalse()) + Expect(registry.Holds(nodeID)).To(BeFalse()) + }) + + It("holds a tunnel and serves its control verbs over it once attached", func() { + Expect(follower.Attach(ctx)).To(Succeed()) + Expect(follower.Attached()).To(BeTrue()) + Eventually(func() bool { return registry.Holds(nodeID) }, "15s").Should(BeTrue()) + + status, body := viaTunnel(registrationToken) + Expect(status).To(Equal(http.StatusOK), body) + Expect(direct()).To(Equal(http.StatusOK)) + }) + + It("refuses a control request over the tunnel that carries no token", func() { + Expect(follower.Attach(ctx)).To(Succeed()) + Eventually(func() bool { return registry.Holds(nodeID) }, "15s").Should(BeTrue()) + status, _ := viaTunnel("") + Expect(status).To(Equal(http.StatusUnauthorized)) + }) + + It("is idempotent: a second attach changes nothing", func() { + Expect(follower.Attach(ctx)).To(Succeed()) + Expect(follower.Attach(ctx)).To(Succeed()) + Eventually(func() bool { return registry.Holds(nodeID) }, "15s").Should(BeTrue()) + }) + + It("lets go of the tunnel and the control path on detach, and can attach again", func() { + Expect(follower.Attach(ctx)).To(Succeed()) + Eventually(func() bool { return registry.Holds(nodeID) }, "15s").Should(BeTrue()) + + Expect(follower.Detach()).To(Succeed()) + Expect(follower.Attached()).To(BeFalse()) + Eventually(func() bool { return registry.Holds(nodeID) }, "15s").Should(BeFalse()) + Expect(direct()).To(Equal(http.StatusNotFound)) + + Expect(follower.Attach(ctx)).To(Succeed(), "the verbs are claimed once and not twice") + Eventually(func() bool { return registry.Holds(nodeID) }, "15s").Should(BeTrue()) + status, _ := viaTunnel(registrationToken) + Expect(status).To(Equal(http.StatusOK)) + }) + + It("refuses to attach a worker that has no frontend to attach to", func() { + f := newTunnelFollower(tunnelFollowerOptions{Config: &Config{}, NodeID: nodeID, HTTPAddr: httpAddr, Token: func() string { return "" }, Slot: slot, Register: func(controlServer) error { return nil }}) + Expect(f.CanAttach()).To(BeFalse()) + Expect(f.Attach(ctx)).To(HaveOccurred()) + Expect(f.Attached()).To(BeFalse()) + }) +}) diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 3d7c86fcc07a..2f624ccd20f4 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -122,6 +122,11 @@ type backendSupervisor struct { // loopback. bindHost string + // follower attaches the worker to a tunnel while it runs. It is nil for a + // worker that cannot attach, and nothing calls it until the worker follows + // the carrier of the cluster. + follower *tunnelFollower + cfg *Config ml *model.ModelLoader systemState *system.SystemState diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 55414f5062b9..e84cc3ad7b75 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -190,9 +190,25 @@ func Run(ctx *cliContext.Context, cfg *Config) error { httpControl *httpControlServer controlMux http.Handler ) - if onTunnel { + // A worker on NATS gets an empty slot instead, when it could attach to a + // tunnel later: the plane is mounted at any time, and until then the path + // answers as it always did on NATS. + var controlSlot *nodes.ControlSlot + follower := newTunnelFollower(tunnelFollowerOptions{ + Config: cfg, + NodeID: nodeID, + HTTPAddr: httpAddr, + Token: tunnelToken, + Reauthorize: reauthorize, + }) + switch { + case onTunnel: httpControl = newHTTPControlServer() controlMux = httpControl + case follower.CanAttach(): + controlSlot = &nodes.ControlSlot{} + follower.slot = controlSlot + controlMux = controlSlot } httpServer, err := nodes.StartFileTransferServerWithControl(httpAddr, stagingDir, cfg.ModelsPath, dataDir, cfg.RegistrationToken, config.DefaultMaxUploadSize, readiness, ephemeralCapacity, controlMux, ml.BackendLogs()) if err != nil { @@ -210,18 +226,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { linkDownMsg string ) if onTunnel { - t, terr := StartTunnel(shutdownCtx, TunnelConfig{ - FrontendURL: cfg.RegisterTo, - NodeID: nodeID, - Token: tunnelToken, - // A frontend that refuses the credential is told who this worker is - // again, so that a rotated credential does not leave it unreachable. - Reauthorize: reauthorize, - // Built by tunnelServices and not inline, so that the routing table, - // which is the security boundary of the tunnel, can be reached from a - // spec without starting a worker. - Services: tunnelServices(cfg, httpAddr), - }) + t, terr := StartTunnel(shutdownCtx, tunnelConfig(cfg, nodeID, httpAddr, tunnelToken, reauthorize)) if terr != nil { nodes.ShutdownFileTransferServer(httpServer) return fmt.Errorf("starting the worker tunnel: %w", terr) @@ -322,23 +327,41 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // The watchdog stops load operations the controller no longer renews. go supervisor.runOperationWatchdog(shutdownCtx) + // registerVerbs claims every verb of the worker on a control server: the + // NATS one, the HTTP one of a worker on the tunnel, or the HTTP one that a + // worker on NATS gains when it attaches to a tunnel. + registerVerbs := func(control controlServer) error { + if err := supervisor.registerLifecycleVerbs(control); err != nil { + return fmt.Errorf("serving the worker lifecycle verbs: %w", err) + } + // Serve the file staging verbs only when S3 is configured + if cfg.StorageURL != "" { + if err := cfg.registerFileStagingVerbs(control, ephemeralCapacity); err != nil { + return fmt.Errorf("serving the file staging verbs: %w", err) + } + } + return nil + } var control controlServer if natsClient != nil { control = newNATSControlServer(natsClient, nodeID) } else { control = httpControl } - if err := supervisor.registerLifecycleVerbs(control); err != nil { + if err := registerVerbs(control); err != nil { nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("serving the worker lifecycle verbs: %w", err) + return err } - - // Serve the file staging verbs only when S3 is configured - if cfg.StorageURL != "" { - if err := cfg.registerFileStagingVerbs(control, ephemeralCapacity); err != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("serving the file staging verbs: %w", err) - } + // The hook of a worker on NATS that attaches to a tunnel later. Nothing calls + // it yet, so a worker on NATS behaves as it did. + if controlSlot != nil { + follower.register = registerVerbs + supervisor.follower = follower + defer func() { + if err := follower.Detach(); err != nil { + xlog.Warn("Closing the worker tunnel failed", "error", err) + } + }() } xlog.Info("Worker ready, waiting for backend.install events") From 13a6f3030da1e227ac14b1a543fba653f93d56fa Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:05:53 +0000 Subject: [PATCH 104/162] fix(application): resync model configs over the broadcaster holder The model config resync from master read the NATS client field that the holders replaced. Pass the broadcaster holder, which forwards the reconnect hooks. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/startup.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/application/startup.go b/core/application/startup.go index f9b27814654a..caaf78f7e83e 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -521,7 +521,7 @@ func New(opts ...config.AppOption) (*Application, error) { // frontend runs its own pass against the shared models directory; the // pass is idempotent, so no leader is needed. if modelConfigResync != nil && distSvc != nil { - modelConfigResync.Start(options.Context, options.Distributed.ModelConfigResyncIntervalOrDefault(), distSvc.Nats) + modelConfigResync.Start(options.Context, options.Distributed.ModelConfigResyncIntervalOrDefault(), distSvc.Broadcaster) } if options.PreloadJSONModels != "" { From 45b9b8234557ee445dd40afa00c6935246201659 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:20:33 +0000 Subject: [PATCH 105/162] feat(cluster): store the settings, the live replicas and the timings of a change of carrier A change of carrier needs three facts that every replica must read the same way. The carrier row now records whether the admin forced the change and a note on how it ended. Its change stamp comes from the clock of the database, and CarrierStore.DBNow reads that clock, so no timeout of a change depends on the clock of a replica. Each replica can report which carriers it could build (Registry.ReportAvailability). A dry run reads this, because a dry run changes no epoch and no replica builds anything for it. Registry.ListLive lists the live replicas with that report. SettingsStore keeps the settings that decide what the whole cluster does (the NATS addresses and the timings of a change) in a table. A runtime setting of the application lives in a file of one process, so two replicas could hold two values. No credential is stored. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/carrier_store.go | 30 ++- .../cluster/carrier_store_extra_test.go | 208 +++++++++++++++++ core/services/cluster/instance.go | 92 ++++++++ core/services/cluster/settings.go | 211 ++++++++++++++++++ 4 files changed, 538 insertions(+), 3 deletions(-) create mode 100644 core/services/cluster/carrier_store_extra_test.go create mode 100644 core/services/cluster/settings.go diff --git a/core/services/cluster/carrier_store.go b/core/services/cluster/carrier_store.go index 73f664703ce5..fb10c043bc2d 100644 --- a/core/services/cluster/carrier_store.go +++ b/core/services/cluster/carrier_store.go @@ -67,10 +67,18 @@ type CarrierRow struct { Draining Carrier `gorm:"size:16;not null;default:''" json:"draining,omitempty"` DrainingUntil *time.Time `json:"draining_until,omitempty"` ChangedBy string `gorm:"size:255" json:"changed_by"` - // ChangedAt is for operators. Nothing decides on it, so replicas need no - // agreed clock. + // ChangedAt is stamped by the database on every change, and the age of a + // change is read on that same clock (CarrierStore.DBNow), so that the + // timeouts of a change need no clock that the replicas agree on. ChangedAt time.Time `json:"changed_at"` PrevEpoch int64 `json:"prev_epoch"` + // Force records that the admin who started the change accepted that a + // replica which is not ready, or a worker that cannot follow, is left + // behind. It belongs to one change and is cleared by the next one. + Force bool `gorm:"not null;default:false" json:"force,omitempty"` + // Note says in a sentence how the last change ended, or why it did. It is + // for operators. + Note string `gorm:"size:512;not null;default:''" json:"note,omitempty"` } func (CarrierRow) TableName() string { return "cluster_carrier" } @@ -84,6 +92,8 @@ type Change struct { Draining Carrier DrainingUntil *time.Time By string + Force bool + Note string } func (c Change) validate() error { @@ -175,7 +185,9 @@ func (s *CarrierStore) Transition(ctx context.Context, from int64, change Change "draining": change.Draining, "draining_until": change.DrainingUntil, "changed_by": change.By, - "changed_at": time.Now().UTC(), + "changed_at": gorm.Expr("now()"), + "force": change.Force, + "note": change.Note, "prev_epoch": gorm.Expr("epoch"), "epoch": gorm.Expr("epoch + 1"), }) @@ -191,3 +203,15 @@ func (s *CarrierStore) Transition(ctx context.Context, from int64, change Change } return row, nil } + +// DBNow returns the time on the clock of the database. The timeouts of a change +// (the prepare timeout, the transition window, the end of a drain) are compared +// with it and with the stamps that the database wrote, never with the clock of a +// replica. +func (s *CarrierStore) DBNow(ctx context.Context) (time.Time, error) { + var now time.Time + if err := s.db.WithContext(ctx).Raw("SELECT now()").Scan(&now).Error; err != nil { + return time.Time{}, fmt.Errorf("reading the clock of the database: %w", err) + } + return now, nil +} diff --git a/core/services/cluster/carrier_store_extra_test.go b/core/services/cluster/carrier_store_extra_test.go new file mode 100644 index 000000000000..a7943d9261f5 --- /dev/null +++ b/core/services/cluster/carrier_store_extra_test.go @@ -0,0 +1,208 @@ +package cluster_test + +import ( + "context" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("The carrier row during a change", func() { + var ( + db *gorm.DB + ctx context.Context + store *cluster.CarrierStore + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + var err error + store, err = cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(ctx, cluster.CarrierNATS, "replica-a") + Expect(err).ToNot(HaveOccurred()) + }) + + It("keeps the force flag and the note of the change that wrote it", func() { + row, err := store.Transition(ctx, 1, cluster.Change{ + Active: cluster.CarrierNATS, State: cluster.StatePrepare, Target: cluster.CarrierTunnel, + Force: true, Note: "requested", By: "admin", + }) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Force).To(BeTrue()) + Expect(row.Note).To(Equal("requested")) + + got, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(got.Force).To(BeTrue()) + Expect(got.Note).To(Equal("requested")) + + row, err = store.Transition(ctx, got.Epoch, cluster.Change{Active: cluster.CarrierNATS, State: cluster.StateStable, By: "leader"}) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Force).To(BeFalse(), "a change writes every field, so the flag does not outlive its change") + Expect(row.Note).To(BeEmpty()) + }) + + It("stamps a change with the clock of the database and reads the age of the row on that clock", func() { + row, err := store.Transition(ctx, 1, cluster.Change{ + Active: cluster.CarrierNATS, State: cluster.StatePrepare, Target: cluster.CarrierTunnel, By: "admin", + }) + Expect(err).ToNot(HaveOccurred()) + + now, err := store.DBNow(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(now.Sub(row.ChangedAt)).To(BeNumerically("~", 0, 5*time.Second)) + Expect(now.Sub(row.ChangedAt)).To(BeNumerically(">=", 0), "the stamp is never ahead of the clock that wrote it") + }) +}) + +var _ = Describe("The live replicas", func() { + var ( + db *gorm.DB + ctx context.Context + reg *cluster.Registry + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("lists the replicas whose heartbeat is inside the window, with what each reported", func() { + Expect(reg.Register(ctx, "a", "v1", 3, "")).To(Succeed()) + Expect(reg.Register(ctx, "b", "v2", 2, "no route to the broker")).To(Succeed()) + Expect(reg.Register(ctx, "gone", "v1", 0, "")).To(Succeed()) + Expect(db.Exec(`UPDATE instances SET last_seen = now() - interval '5 minutes' WHERE id = 'gone'`).Error).To(Succeed()) + + live, err := reg.ListLive(ctx, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(2)) + Expect(live[0].ID).To(Equal("a")) + Expect(live[0].Version).To(Equal("v1")) + Expect(live[0].ReadyEpoch).To(Equal(int64(3))) + Expect(live[1].ID).To(Equal("b")) + Expect(live[1].ReadyReason).To(Equal("no route to the broker")) + }) + + It("records whether a replica can use each carrier and how long ago it looked", func() { + Expect(reg.Register(ctx, "a", "v1", 0, "")).To(Succeed()) + Expect(reg.ReportAvailability(ctx, "a", map[cluster.Carrier]string{ + cluster.CarrierNATS: "no NATS URL is set", + cluster.CarrierTunnel: "", + })).To(Succeed()) + + live, err := reg.ListLive(ctx, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + + reason, known := live[0].AvailabilityFor(cluster.CarrierNATS) + Expect(known).To(BeTrue()) + Expect(reason).To(Equal("no NATS URL is set")) + reason, known = live[0].AvailabilityFor(cluster.CarrierTunnel) + Expect(known).To(BeTrue()) + Expect(reason).To(BeEmpty()) + Expect(live[0].AvailabilityAge).To(BeNumerically("<", 5*time.Second)) + }) + + It("says a carrier is unknown for a replica that never looked", func() { + Expect(reg.Register(ctx, "a", "v1", 0, "")).To(Succeed()) + live, err := reg.ListLive(ctx, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + _, known := live[0].AvailabilityFor(cluster.CarrierNATS) + Expect(known).To(BeFalse()) + }) + + It("keeps what a replica reported when it registers again", func() { + Expect(reg.Register(ctx, "a", "v1", 0, "")).To(Succeed()) + Expect(reg.ReportAvailability(ctx, "a", map[cluster.Carrier]string{cluster.CarrierNATS: ""})).To(Succeed()) + Expect(reg.Register(ctx, "a", "v1", 4, "")).To(Succeed()) + + live, err := reg.ListLive(ctx, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + _, known := live[0].AvailabilityFor(cluster.CarrierNATS) + Expect(known).To(BeTrue()) + }) + + It("refuses to record the availability of a replica that is not registered", func() { + err := reg.ReportAvailability(ctx, "nobody", map[cluster.Carrier]string{cluster.CarrierNATS: ""}) + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + }) +}) + +var _ = Describe("The settings of the cluster", func() { + var ( + ctx context.Context + store *cluster.SettingsStore + ) + + BeforeEach(func() { + db := testutil.SetupTestDB() + ctx = context.Background() + var err error + store, err = cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + }) + + It("reports a setting that was never stored", func() { + _, ok, err := store.Get(ctx, cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeFalse()) + }) + + It("stores a value and replaces it", func() { + Expect(store.Set(ctx, cluster.SettingNATSURL, "nats://one:4222", "admin")).To(Succeed()) + Expect(store.Set(ctx, cluster.SettingNATSURL, "nats://two:4222", "admin")).To(Succeed()) + v, ok, err := store.Get(ctx, cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(v).To(Equal("nats://two:4222")) + }) + + It("inserts a value only when none exists, so that replicas that start together agree", func() { + made, err := store.SetIfAbsent(ctx, cluster.SettingNATSURL, "nats://one:4222", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(made).To(BeTrue()) + made, err = store.SetIfAbsent(ctx, cluster.SettingNATSURL, "nats://two:4222", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(made).To(BeFalse()) + v, _, _ := store.Get(ctx, cluster.SettingNATSURL) + Expect(v).To(Equal("nats://one:4222")) + }) + + It("clears a value", func() { + Expect(store.Set(ctx, cluster.SettingNATSURL, "nats://one:4222", "admin")).To(Succeed()) + Expect(store.Clear(ctx, cluster.SettingNATSURL)).To(Succeed()) + _, ok, err := store.Get(ctx, cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeFalse()) + }) + + It("refuses a key it does not know", func() { + Expect(store.Set(ctx, "nats.password", "secret", "admin")).To(MatchError(cluster.ErrUnknownSetting)) + }) + + It("reads the timings of a change with a default for each one that is not set", func() { + t, err := store.Timings(ctx, cluster.Timings{}) + Expect(err).ToNot(HaveOccurred()) + Expect(t).To(Equal(cluster.DefaultTimings)) + + Expect(store.Set(ctx, cluster.SettingMaxDrain, "20m", "admin")).To(Succeed()) + t, err = store.Timings(ctx, cluster.Timings{PrepareTimeout: 10 * time.Second}) + Expect(err).ToNot(HaveOccurred()) + Expect(t.MaxDrain).To(Equal(20 * time.Minute)) + Expect(t.PrepareTimeout).To(Equal(10*time.Second), "the fallback of the replica applies when the cluster sets nothing") + Expect(t.TransitionWindow).To(Equal(cluster.DefaultTimings.TransitionWindow)) + }) + + It("refuses a timing that is not a positive duration", func() { + Expect(store.Set(ctx, cluster.SettingMaxDrain, "soon", "admin")).To(HaveOccurred()) + Expect(store.Set(ctx, cluster.SettingMaxDrain, "-5s", "admin")).To(HaveOccurred()) + }) +}) diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index 47286fdce1eb..052d42e86c5a 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -2,6 +2,7 @@ package cluster import ( "context" + "encoding/json" "errors" "fmt" "time" @@ -38,6 +39,37 @@ type Instance struct { // secret behind this hash. It is empty for a replica that has not published // one, and an empty hash authorises nobody. PeerTokenHash string `gorm:"size:64" json:"-"` + // Availability says, for each carrier, whether this replica could build it + // the last time it looked: a JSON object from the name of the carrier to the + // reason it cannot, empty when it can. It is what a dry run reads, because a + // dry run changes no epoch and so no replica builds anything for it. + Availability string `gorm:"type:text;not null;default:''" json:"-"` + // AvailabilityAt is when the replica last wrote Availability, by the clock of + // the database. + AvailabilityAt *time.Time `json:"-"` +} + +// LiveInstance is an Instance as ListLive reads it. +type LiveInstance struct { + Instance + // AvailabilityAge is how long ago the replica wrote Availability, on the + // clock of the database. It is zero when the replica never did, which + // AvailabilityFor reports as unknown. + AvailabilityAge time.Duration `gorm:"-" json:"-"` +} + +// AvailabilityFor reports why the replica cannot build carrier, or an empty +// reason when it can. known is false when the replica never reported on it. +func (i LiveInstance) AvailabilityFor(c Carrier) (reason string, known bool) { + if i.Availability == "" { + return "", false + } + var byCarrier map[Carrier]string + if err := json.Unmarshal([]byte(i.Availability), &byCarrier); err != nil { + return "", false + } + reason, known = byCarrier[c] + return reason, known } // Registry reads and writes the instances table and the table of worker @@ -181,3 +213,63 @@ const instanceIsLive = `instances.last_seen > now() - make_interval(secs => ?)` // It is built from instanceIsLive and does not restate it, so the deployment // has one spelling of "alive". const LiveInstanceIDsSQL = `SELECT instances.id FROM instances WHERE ` + instanceIsLive + +// ReportAvailability records, for each carrier, why this replica cannot build it +// (an empty reason when it can). A carrier that the map leaves out keeps what the +// replica reported before. +func (r *Registry) ReportAvailability(ctx context.Context, id string, byCarrier map[Carrier]string) error { + var current Instance + err := r.db.WithContext(ctx).Select("availability").Where("id = ?", id).Take(¤t).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return fmt.Errorf("reporting availability of instance %q: %w", id, ErrInstanceNotFound) + } + if err != nil { + return fmt.Errorf("reporting availability of instance %q: %w", id, err) + } + merged := map[Carrier]string{} + if current.Availability != "" { + _ = json.Unmarshal([]byte(current.Availability), &merged) + } + for c, reason := range byCarrier { + merged[c] = reason + } + raw, err := json.Marshal(merged) + if err != nil { + return fmt.Errorf("encoding availability of instance %q: %w", id, err) + } + res := r.db.WithContext(ctx).Model(&Instance{}).Where("id = ?", id).Updates(map[string]any{ + "availability": string(raw), + "availability_at": gorm.Expr("now()"), + }) + if res.Error != nil { + return fmt.Errorf("reporting availability of instance %q: %w", id, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("reporting availability of instance %q: %w", id, ErrInstanceNotFound) + } + return nil +} + +// ListLive returns the replicas that the deployment considers alive, ordered by +// id. It uses instanceIsLive, the one definition of alive. +func (r *Registry) ListLive(ctx context.Context, within time.Duration) ([]LiveInstance, error) { + type row struct { + Instance + AgeSeconds *float64 `gorm:"column:availability_age"` + } + var rows []row + err := r.db.WithContext(ctx).Raw( + `SELECT instances.*, EXTRACT(EPOCH FROM (now() - instances.availability_at)) AS availability_age + FROM instances WHERE `+instanceIsLive+` ORDER BY instances.id`, within.Seconds()).Scan(&rows).Error + if err != nil { + return nil, fmt.Errorf("listing the live instances: %w", err) + } + out := make([]LiveInstance, len(rows)) + for i, rw := range rows { + out[i] = LiveInstance{Instance: rw.Instance} + if rw.AgeSeconds != nil { + out[i].AvailabilityAge = time.Duration(*rw.AgeSeconds * float64(time.Second)) + } + } + return out, nil +} diff --git a/core/services/cluster/settings.go b/core/services/cluster/settings.go new file mode 100644 index 000000000000..af6223759a25 --- /dev/null +++ b/core/services/cluster/settings.go @@ -0,0 +1,211 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// Names of the settings that every replica shares through the database. +// +// A runtime setting of the application lives in a file of one process, so two +// replicas can hold two values. These settings decide what the whole cluster +// does, so they live in a table. No credential is stored here: NATS keys and +// TLS files stay on each replica, as they always did. +const ( + // SettingNATSURL is the address of the NATS server that the frontends use. + SettingNATSURL = "nats.url" + // SettingNATSWorkerURL is the address that workers are told to use, when it + // differs from the one the frontends use. It defaults to SettingNATSURL. + SettingNATSWorkerURL = "nats.worker_url" + // SettingPrepareTimeout is how long a change of carrier waits for every + // replica to be ready. + SettingPrepareTimeout = "switch.prepare_timeout" + // SettingTransitionWindow is how long a change waits for every replica to + // confirm that it uses the new carrier. + SettingTransitionWindow = "switch.transition_window" + // SettingMaxDrain is how long the previous carrier stays attached after a + // change. + SettingMaxDrain = "switch.max_drain" +) + +// ErrUnknownSetting means a key is not one of the cluster settings. +var ErrUnknownSetting = errors.New("unknown cluster setting") + +var knownSettings = map[string]func(string) error{ + SettingNATSURL: nil, + SettingNATSWorkerURL: nil, + SettingPrepareTimeout: durationSetting, + SettingTransitionWindow: durationSetting, + SettingMaxDrain: durationSetting, +} + +func durationSetting(v string) error { + d, err := time.ParseDuration(v) + if err != nil { + return fmt.Errorf("not a duration: %w", err) + } + if d <= 0 { + return errors.New("a duration must be positive") + } + return nil +} + +// Timings are the waits of a change. +type Timings struct { + // PrepareTimeout is how long prepare waits for every live replica. + PrepareTimeout time.Duration + // TransitionWindow is how long commit waits for every live replica to + // confirm. + TransitionWindow time.Duration + // MaxDrain is how long the previous carrier stays attached after the commit. + // Work that is still running on it then is failed by the reaper. + MaxDrain time.Duration +} + +// DefaultTimings are used for a timing that nothing sets. +var DefaultTimings = Timings{ + PrepareTimeout: time.Minute, + TransitionWindow: 2 * time.Minute, + MaxDrain: 15 * time.Minute, +} + +func (t Timings) withDefaults() Timings { + if t.PrepareTimeout <= 0 { + t.PrepareTimeout = DefaultTimings.PrepareTimeout + } + if t.TransitionWindow <= 0 { + t.TransitionWindow = DefaultTimings.TransitionWindow + } + if t.MaxDrain <= 0 { + t.MaxDrain = DefaultTimings.MaxDrain + } + return t +} + +// Setting is one row of table cluster_settings. +type Setting struct { + Key string `gorm:"primaryKey;size:64"` + Value string `gorm:"type:text;not null"` + UpdatedBy string `gorm:"size:255"` + UpdatedAt time.Time `gorm:"autoUpdateTime"` +} + +func (Setting) TableName() string { return "cluster_settings" } + +// SettingsStore reads and writes the cluster settings. +type SettingsStore struct{ db *gorm.DB } + +// NewSettingsStore creates the table when it is missing. +func NewSettingsStore(db *gorm.DB) (*SettingsStore, error) { + if err := db.AutoMigrate(&Setting{}); err != nil { + return nil, fmt.Errorf("migrating cluster_settings: %w", err) + } + return &SettingsStore{db: db}, nil +} + +func checkSetting(key, value string) error { + validate, ok := knownSettings[key] + if !ok { + return fmt.Errorf("%w: %q", ErrUnknownSetting, key) + } + if validate != nil { + if err := validate(value); err != nil { + return fmt.Errorf("setting %s: %w", key, err) + } + } + return nil +} + +// Get returns the value of a setting, and false when it was never stored. +func (s *SettingsStore) Get(ctx context.Context, key string) (string, bool, error) { + var row Setting + err := s.db.WithContext(ctx).Where("key = ?", key).Take(&row).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return "", false, nil + } + if err != nil { + return "", false, fmt.Errorf("reading cluster setting %s: %w", key, err) + } + return row.Value, true, nil +} + +// Set stores a value, replacing the one that was there. +func (s *SettingsStore) Set(ctx context.Context, key, value, by string) error { + if err := checkSetting(key, value); err != nil { + return err + } + err := s.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "key"}}, + DoUpdates: clause.Assignments(map[string]any{"value": value, "updated_by": by, "updated_at": gorm.Expr("now()")}), + }).Create(&Setting{Key: key, Value: value, UpdatedBy: by}).Error + if err != nil { + return fmt.Errorf("storing cluster setting %s: %w", key, err) + } + return nil +} + +// SetIfAbsent stores a value only when the key has none, and reports whether it +// did. Replicas that start together can all call it: the first insert wins. +func (s *SettingsStore) SetIfAbsent(ctx context.Context, key, value, by string) (bool, error) { + if err := checkSetting(key, value); err != nil { + return false, err + } + res := s.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}). + Create(&Setting{Key: key, Value: value, UpdatedBy: by}) + if res.Error != nil { + return false, fmt.Errorf("storing cluster setting %s: %w", key, res.Error) + } + return res.RowsAffected == 1, nil +} + +// Clear removes a setting. +func (s *SettingsStore) Clear(ctx context.Context, key string) error { + if _, ok := knownSettings[key]; !ok { + return fmt.Errorf("%w: %q", ErrUnknownSetting, key) + } + if err := s.db.WithContext(ctx).Where("key = ?", key).Delete(&Setting{}).Error; err != nil { + return fmt.Errorf("clearing cluster setting %s: %w", key, err) + } + return nil +} + +// All returns every stored setting by key. +func (s *SettingsStore) All(ctx context.Context) (map[string]string, error) { + var rows []Setting + if err := s.db.WithContext(ctx).Find(&rows).Error; err != nil { + return nil, fmt.Errorf("reading cluster settings: %w", err) + } + out := make(map[string]string, len(rows)) + for _, r := range rows { + out[r.Key] = r.Value + } + return out, nil +} + +// Timings returns the timings of a change. A timing that the cluster did not set +// comes from fallback, and one that fallback leaves at zero comes from +// DefaultTimings. +func (s *SettingsStore) Timings(ctx context.Context, fallback Timings) (Timings, error) { + all, err := s.All(ctx) + if err != nil { + return Timings{}, err + } + t := fallback + for key, dst := range map[string]*time.Duration{ + SettingPrepareTimeout: &t.PrepareTimeout, + SettingTransitionWindow: &t.TransitionWindow, + SettingMaxDrain: &t.MaxDrain, + } { + if v, ok := all[key]; ok { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + *dst = d + } + } + } + return t.withDefaults(), nil +} From 5b388be2dfb1e9b5a87eb29043b4db525ded2531 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:20:33 +0000 Subject: [PATCH 106/162] feat(cluster): add the protocol of a change of carrier Switch moves the carrier row through stable, prepare, commit and stable, and from prepare back to stable on an abort. Every move is a compare-and-set on the epoch, and the protocol keeps no state outside the row and the instances table. A leader that dies between two moves is replaced by one that reads the same row. Preflight reports what a change would meet: the live replicas with their versions, the replicas that cannot build the target, the workers that cannot follow and why, and the work in flight. A request goes ahead only when nothing blocks. A forced request goes ahead past a replica or a worker that is not ready, and says so in the row. Drive makes the next move that is due. It commits when every live replica is ready for the epoch of the prepare. It aborts when a replica says it cannot build the target or does not become ready in time, and it names the replica. A forced change commits without it. It settles when every live replica confirms the commit, or when the transition window ends. It ends the drain at the time the commit set. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/switch.go | 568 +++++++++++++++++++++++++++ core/services/cluster/switch_test.go | 561 ++++++++++++++++++++++++++ 2 files changed, 1129 insertions(+) create mode 100644 core/services/cluster/switch.go create mode 100644 core/services/cluster/switch_test.go diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go new file mode 100644 index 000000000000..83cedab8457f --- /dev/null +++ b/core/services/cluster/switch.go @@ -0,0 +1,568 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "slices" + "strings" + "time" + + "github.com/mudler/xlog" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// This file is the protocol of a change of carrier. It moves the cluster row +// through +// +// stable -> prepare -> commit -> stable +// +// and from prepare back to stable (abort). Every move is a compare-and-set on +// the epoch of the row (CarrierStore.Transition), so two admins, or two +// replicas that both think they lead, cannot interleave. The protocol keeps no +// state outside the row and the instances table: a leader that dies between two +// moves is replaced by one that reads the same row and goes on. Every timeout is +// measured on the clock of the database. +// +// The replicas do the work of a change. In prepare each one builds the target +// carrier, listens on it, and reports its epoch as ready. In commit each one +// starts to publish on the target and confirms with the epoch of the commit. +// The leader (Switch.Drive) only decides when the next move is due. + +var ( + // ErrBusy means a change is under way, or another writer started one first. + ErrBusy = errors.New("a change of carrier is already under way") + // ErrNotAbortable means there is nothing to abort. Only prepare can be + // aborted: after the commit the target is the active carrier, and going back + // is a change like any other. + ErrNotAbortable = errors.New("only a change that is being prepared can be aborted") +) + +// WorkerInfo is what the switch needs to know about one worker. +type WorkerInfo struct { + ID string + Name string + // Type is the node type: backend or agent. + Type string + // Attached lists the carriers the worker is connected to now. + Attached []Carrier + // Reports is true when the worker reports which carriers it can follow. A + // worker that predates carrier switching reports nothing. + Reports bool + // Follow lists the carriers the worker can attach to. + Follow []Carrier + // FollowError is why the worker cannot follow, as it reported. + FollowError string +} + +func (w WorkerInfo) canFollow(target Carrier) (bool, string) { + if slices.Contains(w.Attached, target) { + return true, "" + } + if !w.Reports { + return false, "the worker predates carrier switching: it reports no capabilities and stays on its carrier" + } + if slices.Contains(w.Follow, target) { + return true, "" + } + reason := fmt.Sprintf("the worker cannot use the %s carrier", target) + if w.FollowError != "" { + reason += ": " + w.FollowError + } + return false, reason +} + +// WorkerSource lists the workers that are alive. +type WorkerSource interface { + Workers(ctx context.Context) ([]WorkerInfo, error) +} + +// InFlight counts the work that a change of carrier can touch. It is +// information for the admin, and it never blocks a change. +type InFlight struct { + // Loads is the number of models that are loading. + Loads int `json:"loads"` + // Jobs is the number of jobs that run. + Jobs int `json:"jobs"` + // PendingClaims and ClaimedClaims count the rows of the claim queue. + PendingClaims int `json:"pending_claims"` + ClaimedClaims int `json:"claimed_claims"` +} + +// WorkSource counts the work in flight. +type WorkSource interface { + InFlight(ctx context.Context) (InFlight, error) +} + +// Kinds of Blocker. +const ( + BlockerBusy = "busy" + BlockerSame = "same" + BlockerReplica = "replica" + BlockerWorker = "worker" +) + +// Blocker is one reason a change would not go ahead. +type Blocker struct { + Kind string `json:"kind"` + ID string `json:"id,omitempty"` + Reason string `json:"reason"` + // Forceable is true when a forced request goes ahead anyway. + Forceable bool `json:"forceable"` +} + +// ReplicaStatus is one live frontend replica. +type ReplicaStatus struct { + ID string `json:"id"` + Version string `json:"version"` + ReadyEpoch int64 `json:"ready_epoch"` + ReadyReason string `json:"ready_reason,omitempty"` + // Availability maps each carrier to the reason the replica cannot build it, + // empty when it can. A carrier it never looked at is absent. + Availability map[Carrier]string `json:"availability,omitempty"` + AvailabilityAge time.Duration `json:"availability_age_ns,omitempty"` +} + +// WorkerStatus is one worker, and whether it can follow the target. +type WorkerStatus struct { + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` + Attached []Carrier `json:"attached,omitempty"` + Follow []Carrier `json:"follow,omitempty"` + FollowError string `json:"follow_error,omitempty"` + // CanFollow and Reason are set when the report has a target. + CanFollow bool `json:"can_follow"` + Reason string `json:"reason,omitempty"` +} + +// Report is the state of the cluster for the admin. With a target it is also the +// answer of a dry run. +type Report struct { + Row CarrierRow `json:"row"` + Active Carrier `json:"active"` + Target Carrier `json:"target,omitempty"` + Epoch int64 `json:"epoch"` + State State `json:"state"` + OK bool `json:"ok"` + Blockers []Blocker `json:"blockers,omitempty"` + Warnings []string `json:"warnings,omitempty"` + Replicas []ReplicaStatus `json:"replicas"` + Workers []WorkerStatus `json:"workers"` + InFlight InFlight `json:"in_flight"` + Timings Timings `json:"timings"` + // DrainRemaining is how long the previous carrier stays attached, when one + // is draining. + DrainRemaining time.Duration `json:"drain_remaining_ns,omitempty"` +} + +// Request asks for a change of carrier. +type Request struct { + Target Carrier + // By names who asked. It goes in the row and in the log. + By string + // Force goes ahead past a replica or a worker that is not ready. + Force bool +} + +// BlockedError is returned by Switch.Request when the preflight blocks it. The +// row is unchanged. +type BlockedError struct{ Report Report } + +func (e *BlockedError) Error() string { + parts := make([]string, 0, len(e.Report.Blockers)) + for _, b := range e.Report.Blockers { + if b.ID != "" { + parts = append(parts, b.ID+": "+b.Reason) + } else { + parts = append(parts, b.Reason) + } + } + return "the change of carrier is blocked: " + strings.Join(parts, "; ") +} + +// SwitchOptions are the inputs of NewSwitch. +type SwitchOptions struct { + Store *CarrierStore + Registry *Registry + // Workers and Work may be nil; the report is then empty on those points. + Workers WorkerSource + Work WorkSource + // Timings returns the waits of a change, read at each use so that a change of + // the settings applies to the next decision. + Timings func() Timings + // Liveness is the window of replica liveness. InstanceLiveness when zero. + Liveness time.Duration + // AvailabilityMaxAge is how old a replica's report of what it can build may + // be. Two minutes when zero. + AvailabilityMaxAge time.Duration + Meter metric.Meter +} + +// Switch is the protocol of a change of carrier. +type Switch struct { + o SwitchOptions + changes metric.Int64Counter + phase metric.Float64Histogram +} + +// NewSwitch returns the protocol over the row and the registry. +func NewSwitch(o SwitchOptions) (*Switch, error) { + if o.Store == nil || o.Registry == nil { + return nil, errors.New("a switch needs the carrier store and the instance registry") + } + if o.Liveness <= 0 { + o.Liveness = InstanceLiveness + } + if o.AvailabilityMaxAge <= 0 { + o.AvailabilityMaxAge = 2 * time.Minute + } + if o.Timings == nil { + o.Timings = func() Timings { return DefaultTimings } + } + meter := o.Meter + if meter == nil { + meter = otel.Meter("github.com/mudler/LocalAI") + } + s := &Switch{o: o} + // An instrument that cannot be created is left nil: a metric must not stop a + // change of carrier. + s.changes, _ = meter.Int64Counter("localai_carrier_changes_total", + metric.WithDescription("Moves of the cluster carrier row made by the leader or an admin, by outcome")) + s.phase, _ = meter.Float64Histogram("localai_carrier_phase_duration_seconds", + metric.WithDescription("How long the cluster stayed in prepare and in commit"), metric.WithUnit("s")) + return s, nil +} + +func (s *Switch) count(outcome string) { + if s.changes != nil { + s.changes.Add(context.Background(), 1, metric.WithAttributes(attribute.String("outcome", outcome))) + } +} + +func (s *Switch) observePhase(phase string, d time.Duration) { + if s.phase != nil { + s.phase.Record(context.Background(), d.Seconds(), metric.WithAttributes(attribute.String("phase", phase))) + } +} + +func other(c Carrier) Carrier { + if c == CarrierNATS { + return CarrierTunnel + } + return CarrierNATS +} + +// Status reports the cluster as it is, with no target. +func (s *Switch) Status(ctx context.Context) (Report, error) { + return s.report(ctx, "") +} + +// Preflight reports what a change to target would meet. It changes nothing. +func (s *Switch) Preflight(ctx context.Context, target Carrier) (Report, error) { + if !target.valid() { + return Report{}, fmt.Errorf("%w: %q", ErrInvalidCarrier, target) + } + return s.report(ctx, target) +} + +func (s *Switch) report(ctx context.Context, target Carrier) (Report, error) { + row, err := s.o.Store.Get(ctx) + if err != nil { + return Report{}, err + } + r := Report{Row: row, Active: row.Active, Target: target, Epoch: row.Epoch, State: row.State, Timings: s.o.Timings().withDefaults()} + + if row.Draining != "" && row.DrainingUntil != nil { + if now, err := s.o.Store.DBNow(ctx); err == nil && now.Before(*row.DrainingUntil) { + r.DrainRemaining = row.DrainingUntil.Sub(now) + } + } + + live, err := s.o.Registry.ListLive(ctx, s.o.Liveness) + if err != nil { + return Report{}, err + } + for _, in := range live { + rs := ReplicaStatus{ID: in.ID, Version: in.Version, ReadyEpoch: in.ReadyEpoch, ReadyReason: in.ReadyReason, AvailabilityAge: in.AvailabilityAge} + for _, c := range []Carrier{CarrierNATS, CarrierTunnel} { + if reason, known := in.AvailabilityFor(c); known { + if rs.Availability == nil { + rs.Availability = map[Carrier]string{} + } + rs.Availability[c] = reason + } + } + r.Replicas = append(r.Replicas, rs) + } + + var workers []WorkerInfo + if s.o.Workers != nil { + if workers, err = s.o.Workers.Workers(ctx); err != nil { + return Report{}, err + } + } + for _, w := range workers { + ws := WorkerStatus{ID: w.ID, Name: w.Name, Type: w.Type, Attached: w.Attached, Follow: w.Follow, FollowError: w.FollowError} + if target != "" { + ws.CanFollow, ws.Reason = w.canFollow(target) + } + r.Workers = append(r.Workers, ws) + } + if s.o.Work != nil { + if r.InFlight, err = s.o.Work.InFlight(ctx); err != nil { + return Report{}, err + } + } + + if target == "" { + return r, nil + } + switch { + case row.State != StateStable: + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerBusy, Reason: fmt.Sprintf("a change to %s is under way (%s); only an abort is accepted", row.Target, row.State)}) + case target == row.Active: + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerSame, Reason: fmt.Sprintf("%s is the active carrier already", target)}) + } + if len(r.Blockers) == 0 { + for _, in := range live { + reason, known := in.AvailabilityFor(target) + switch { + case !known: + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerReplica, ID: in.ID, Forceable: true, + Reason: fmt.Sprintf("the replica has not reported whether it can use the %s carrier", target)}) + case in.AvailabilityAge > s.o.AvailabilityMaxAge: + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerReplica, ID: in.ID, Forceable: true, + Reason: fmt.Sprintf("what the replica reported about the %s carrier is stale (%s old)", target, in.AvailabilityAge.Round(time.Second))}) + case reason != "": + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerReplica, ID: in.ID, Forceable: true, + Reason: fmt.Sprintf("the replica cannot use the %s carrier: %s", target, reason)}) + } + } + for _, w := range r.Workers { + if !w.CanFollow { + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerWorker, ID: w.ID, Forceable: true, Reason: w.Reason}) + } + } + if row.Draining != "" { + r.Warnings = append(r.Warnings, fmt.Sprintf("the %s carrier is still draining from the last change; this change starts a second drain on top of it", row.Draining)) + } + } + r.OK = len(r.Blockers) == 0 + return r, nil +} + +// Request starts a change to req.Target: it runs the preflight and, when nothing +// blocks, moves the row to prepare. A blocker that can be forced does not block a +// forced request. The row is unchanged when the request is refused. +func (s *Switch) Request(ctx context.Context, req Request) (CarrierRow, Report, error) { + report, err := s.Preflight(ctx, req.Target) + if err != nil { + return CarrierRow{}, Report{}, err + } + var blocking []Blocker + for _, b := range report.Blockers { + if b.Kind == BlockerBusy { + return CarrierRow{}, report, ErrBusy + } + if !b.Forceable || !req.Force { + blocking = append(blocking, b) + } + } + if len(blocking) > 0 { + report.Blockers = blocking + return CarrierRow{}, report, &BlockedError{Report: report} + } + row := report.Row + note := "" + if req.Force && len(report.Blockers) > 0 { + note = fmt.Sprintf("forced past %d blocker(s)", len(report.Blockers)) + } + next, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + Active: row.Active, State: StatePrepare, Target: req.Target, + Draining: row.Draining, DrainingUntil: row.DrainingUntil, + Force: req.Force, Note: note, By: req.By, + }) + if errors.Is(err, ErrStaleEpoch) { + return CarrierRow{}, report, ErrBusy + } + if err != nil { + return CarrierRow{}, report, err + } + s.count("requested") + xlog.Info("Change of carrier requested", "by", req.By, "from", row.Active, "to", req.Target, "epoch", next.Epoch, "force", req.Force) + return next, report, nil +} + +// Abort returns a cluster that is in prepare to stable on the carrier it had. +func (s *Switch) Abort(ctx context.Context, by string) (CarrierRow, error) { + row, err := s.o.Store.Get(ctx) + if err != nil { + return CarrierRow{}, err + } + if row.State != StatePrepare { + return CarrierRow{}, ErrNotAbortable + } + next, err := s.abort(ctx, row, by, fmt.Sprintf("aborted by %s", by)) + if errors.Is(err, ErrStaleEpoch) { + // The leader moved the row first. Say what it is now. + return CarrierRow{}, ErrNotAbortable + } + return next, err +} + +func (s *Switch) abort(ctx context.Context, row CarrierRow, by, note string) (CarrierRow, error) { + next, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + Active: row.Active, State: StateStable, + Draining: row.Draining, DrainingUntil: row.DrainingUntil, + Note: note, By: by, + }) + if err != nil { + return CarrierRow{}, err + } + s.count("aborted") + xlog.Warn("Change of carrier aborted", "to", row.Target, "kept", row.Active, "by", by, "reason", note, "epoch", next.Epoch) + return next, nil +} + +// Drive makes the next move that is due, if any. The leader calls it on a tick. +// It is safe to call from any number of replicas: each move is a compare-and-set +// on the epoch that it read, so a second caller finds the row moved and changes +// nothing. A caller that finds the row moved is not in error. +func (s *Switch) Drive(ctx context.Context) error { + row, err := s.o.Store.Get(ctx) + if err != nil { + return err + } + now, err := s.o.Store.DBNow(ctx) + if err != nil { + return err + } + timings := s.o.Timings().withDefaults() + age := now.Sub(row.ChangedAt) + + switch row.State { + case StatePrepare: + err = s.drivePrepare(ctx, row, now, age, timings) + case StateCommit: + err = s.driveCommit(ctx, row, age, timings) + case StateStable: + err = s.driveDrain(ctx, row, now) + } + if errors.Is(err, ErrStaleEpoch) { + return nil + } + return err +} + +func (s *Switch) drivePrepare(ctx context.Context, row CarrierRow, now time.Time, age time.Duration, t Timings) error { + live, err := s.o.Registry.ListLive(ctx, s.o.Liveness) + if err != nil { + return err + } + var pending, failed []string + reasons := map[string]string{} + for _, in := range live { + switch { + case in.ReadyEpoch != row.Epoch: + pending = append(pending, in.ID) + case in.ReadyReason != "": + failed = append(failed, in.ID) + reasons[in.ID] = in.ReadyReason + } + } + describe := func(ids []string) string { + out := make([]string, len(ids)) + for i, id := range ids { + out[i] = id + if r := reasons[id]; r != "" { + out[i] += " (" + r + ")" + } + } + return strings.Join(out, ", ") + } + + commit := func(note string) error { + until := now.Add(t.MaxDrain) + _, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + Active: row.Target, State: StateCommit, Target: row.Target, + Draining: row.Active, DrainingUntil: &until, + Force: row.Force, Note: note, By: row.ChangedBy, + }) + if err == nil { + s.count("committed") + s.observePhase("prepare", age) + xlog.Info("Change of carrier committed", "to", row.Target, "from", row.Active, "prepared_for", age.Round(time.Millisecond), "note", note) + } + return err + } + abort := func(note string) error { + _, err := s.abort(ctx, row, row.ChangedBy, note) + return err + } + + switch { + case len(failed) > 0 && !row.Force: + return abort("aborted: the replica cannot build the target: " + describe(failed)) + case len(pending) == 0: + if len(failed) > 0 { + return commit("committed without the replica that cannot build the target: " + describe(failed)) + } + return commit("") + case age >= t.PrepareTimeout && row.Force: + return commit("committed without the replica that is not ready: " + describe(append(pending, failed...))) + case age >= t.PrepareTimeout: + return abort(fmt.Sprintf("aborted: the replica was not ready after %s: %s", t.PrepareTimeout, describe(append(pending, failed...)))) + } + return nil +} + +func (s *Switch) driveCommit(ctx context.Context, row CarrierRow, age time.Duration, t Timings) error { + live, err := s.o.Registry.ListLive(ctx, s.o.Liveness) + if err != nil { + return err + } + var missing []string + for _, in := range live { + if in.ReadyEpoch != row.Epoch || in.ReadyReason != "" { + missing = append(missing, in.ID) + } + } + if len(missing) > 0 && age < t.TransitionWindow { + return nil + } + note := row.Note + if len(missing) > 0 { + note = strings.TrimSpace(note + " The replica did not confirm the commit within " + t.TransitionWindow.String() + ": " + strings.Join(missing, ", ")) + } + _, err = s.o.Store.Transition(ctx, row.Epoch, Change{ + Active: row.Active, State: StateStable, + Draining: row.Draining, DrainingUntil: row.DrainingUntil, + Force: row.Force, Note: note, By: row.ChangedBy, + }) + if err == nil { + s.count("settled") + s.observePhase("commit", age) + xlog.Info("Change of carrier settled", "active", row.Active, "draining", row.Draining, "epoch", row.Epoch+1, "unconfirmed", missing) + } + return err +} + +func (s *Switch) driveDrain(ctx context.Context, row CarrierRow, now time.Time) error { + if row.Draining == "" { + return nil + } + if row.DrainingUntil != nil && now.Before(*row.DrainingUntil) { + return nil + } + _, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + Active: row.Active, State: StateStable, Note: "drained", By: row.ChangedBy, + }) + if err == nil { + s.count("drained") + xlog.Info("The previous carrier is released", "released", row.Draining, "active", row.Active, "epoch", row.Epoch+1) + } + return err +} diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go new file mode 100644 index 000000000000..00adedc892a4 --- /dev/null +++ b/core/services/cluster/switch_test.go @@ -0,0 +1,561 @@ +package cluster_test + +import ( + "context" + "errors" + "slices" + "sync" + "sync/atomic" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// fakeWorkers is the list of workers that the switch reads. +type fakeWorkers struct { + mu sync.Mutex + workers []cluster.WorkerInfo +} + +func (f *fakeWorkers) Workers(context.Context) ([]cluster.WorkerInfo, error) { + f.mu.Lock() + defer f.mu.Unlock() + return slices.Clone(f.workers), nil +} + +func (f *fakeWorkers) set(w ...cluster.WorkerInfo) { + f.mu.Lock() + defer f.mu.Unlock() + f.workers = w +} + +type fakeWork struct{ inflight cluster.InFlight } + +func (f *fakeWork) InFlight(context.Context) (cluster.InFlight, error) { return f.inflight, nil } + +var _ = Describe("The switch of the carrier", func() { + var ( + db *gorm.DB + ctx context.Context + store *cluster.CarrierStore + reg *cluster.Registry + workers *fakeWorkers + work *fakeWork + sw *cluster.Switch + timings cluster.Timings + ) + + // replica registers a live replica. It reports the carriers it can build. + replica := func(id string, canTunnel, canNATS string) { + GinkgoHelper() + Expect(reg.Register(ctx, id, "v-"+id, 0, "")).To(Succeed()) + Expect(reg.ReportAvailability(ctx, id, map[cluster.Carrier]string{ + cluster.CarrierTunnel: canTunnel, + cluster.CarrierNATS: canNATS, + })).To(Succeed()) + } + ready := func(id string, epoch int64, reason string) { + GinkgoHelper() + Expect(reg.ReportReady(ctx, id, epoch, reason)).To(Succeed()) + } + row := func() cluster.CarrierRow { + GinkgoHelper() + r, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + return r + } + request := func(target cluster.Carrier, force bool) (cluster.CarrierRow, cluster.Report, error) { + return sw.Request(ctx, cluster.Request{Target: target, By: "admin", Force: force}) + } + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + var err error + store, err = cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(ctx, cluster.CarrierNATS, "replica-a") + Expect(err).ToNot(HaveOccurred()) + reg = cluster.NewRegistry(db) + workers, work = &fakeWorkers{}, &fakeWork{} + timings = cluster.Timings{PrepareTimeout: time.Minute, TransitionWindow: time.Minute, MaxDrain: time.Minute} + sw, err = cluster.NewSwitch(cluster.SwitchOptions{ + Store: store, Registry: reg, Workers: workers, Work: work, + Timings: func() cluster.Timings { return timings }, + }) + Expect(err).ToNot(HaveOccurred()) + }) + + Describe("the preflight", func() { + It("passes when every live replica can build the target and every worker can follow", func() { + replica("a", "", "") + replica("b", "", "") + workers.set(cluster.WorkerInfo{ID: "w1", Name: "w1", Attached: []cluster.Carrier{cluster.CarrierNATS}, Reports: true, + Follow: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}}) + + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(BeEmpty()) + Expect(report.OK).To(BeTrue()) + Expect(report.Active).To(Equal(cluster.CarrierNATS)) + Expect(report.Target).To(Equal(cluster.CarrierTunnel)) + }) + + It("lists the live replicas with their versions, so the admin can check the list is complete", func() { + replica("a", "", "") + replica("b", "", "") + Expect(reg.Register(ctx, "dead", "v-dead", 0, "")).To(Succeed()) + Expect(db.Exec(`UPDATE instances SET last_seen = now() - interval '10 minutes' WHERE id = 'dead'`).Error).To(Succeed()) + + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + var ids []string + for _, r := range report.Replicas { + ids = append(ids, r.ID) + Expect(r.Version).To(Equal("v-" + r.ID)) + } + Expect(ids).To(Equal([]string{"a", "b"})) + }) + + It("blocks on a replica that never said whether it can build the target", func() { + Expect(reg.Register(ctx, "a", "v1", 0, "")).To(Succeed()) + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.OK).To(BeFalse()) + Expect(report.Blockers).To(ConsistOf(HaveField("ID", "a"))) + Expect(report.Blockers[0].Kind).To(Equal(cluster.BlockerReplica)) + Expect(report.Blockers[0].Reason).To(ContainSubstring("has not reported")) + }) + + It("blocks on a replica that says why it cannot build the target", func() { + replica("a", "", "") + replica("b", "LISTEN is refused by the database", "") + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(ConsistOf(SatisfyAll( + HaveField("ID", "b"), + HaveField("Reason", ContainSubstring("LISTEN is refused by the database")), + ))) + }) + + It("ignores what a replica reported long ago", func() { + replica("a", "", "") + Expect(db.Exec(`UPDATE instances SET availability_at = now() - interval '1 hour' WHERE id = 'a'`).Error).To(Succeed()) + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(ConsistOf(HaveField("Reason", ContainSubstring("stale")))) + }) + + It("blocks on a worker that reports no capabilities, and says why", func() { + replica("a", "", "") + workers.set(cluster.WorkerInfo{ID: "old", Name: "old-worker", Attached: []cluster.Carrier{cluster.CarrierNATS}}) + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.OK).To(BeFalse()) + Expect(report.Blockers).To(ConsistOf(SatisfyAll( + HaveField("Kind", cluster.BlockerWorker), + HaveField("ID", "old"), + HaveField("Reason", ContainSubstring("predates")), + ))) + Expect(report.Workers).To(ConsistOf(SatisfyAll(HaveField("ID", "old"), HaveField("CanFollow", false)))) + }) + + It("blocks on a worker that cannot follow, with the reason it reported", func() { + replica("a", "", "") + workers.set(cluster.WorkerInfo{ID: "w", Name: "nats-only", Reports: true, Attached: []cluster.Carrier{cluster.CarrierNATS}, + Follow: []cluster.Carrier{cluster.CarrierNATS}, FollowError: "has no tunnel credential"}) + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(ConsistOf(SatisfyAll( + HaveField("ID", "w"), + HaveField("Reason", ContainSubstring("has no tunnel credential")), + ))) + }) + + It("does not block on a worker that is attached to the target already", func() { + replica("a", "", "") + workers.set(cluster.WorkerInfo{ID: "w", Name: "w", Attached: []cluster.Carrier{cluster.CarrierTunnel}}) + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(BeEmpty()) + }) + + It("reports what is in flight, and does not block on it", func() { + replica("a", "", "") + work.inflight = cluster.InFlight{Loads: 2, Jobs: 3, PendingClaims: 4, ClaimedClaims: 1} + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.OK).To(BeTrue()) + Expect(report.InFlight).To(Equal(work.inflight)) + }) + + It("refuses a target that is already active", func() { + replica("a", "", "") + report, err := sw.Preflight(ctx, cluster.CarrierNATS) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(ConsistOf(HaveField("Kind", cluster.BlockerSame))) + Expect(report.Blockers[0].Forceable).To(BeFalse()) + }) + + It("refuses a carrier that the cluster does not know", func() { + _, err := sw.Preflight(ctx, cluster.Carrier("pigeon")) + Expect(err).To(MatchError(cluster.ErrInvalidCarrier)) + }) + + It("says the cluster is busy while a change is under way", func() { + replica("a", "", "") + _, _, err := request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + report, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.Blockers).To(ConsistOf(HaveField("Kind", cluster.BlockerBusy))) + }) + + It("warns that two drains overlap when a change starts before the last one has drained", func() { + replica("a", "", "") + _, _, err := request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) // commit + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) // stable, draining + Expect(row().Draining).To(Equal(cluster.CarrierNATS)) + + report, err := sw.Preflight(ctx, cluster.CarrierNATS) + Expect(err).ToNot(HaveOccurred()) + Expect(report.OK).To(BeTrue()) + Expect(report.Warnings).To(ContainElement(ContainSubstring("still draining"))) + }) + }) + + Describe("a request", func() { + It("moves a stable cluster to prepare and names the target", func() { + replica("a", "", "") + got, report, err := request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + Expect(report.OK).To(BeTrue()) + Expect(got.State).To(Equal(cluster.StatePrepare)) + Expect(got.Target).To(Equal(cluster.CarrierTunnel)) + Expect(got.Active).To(Equal(cluster.CarrierNATS)) + Expect(got.Epoch).To(Equal(int64(2))) + Expect(got.ChangedBy).To(Equal("admin")) + Expect(got.Force).To(BeFalse()) + }) + + It("changes nothing and returns the report when the preflight blocks", func() { + replica("a", "", "") + workers.set(cluster.WorkerInfo{ID: "old", Name: "old"}) + _, _, err := request(cluster.CarrierTunnel, false) + var blocked *cluster.BlockedError + Expect(errors.As(err, &blocked)).To(BeTrue()) + Expect(blocked.Report.Blockers).ToNot(BeEmpty()) + Expect(row().Epoch).To(Equal(int64(1))) + Expect(row().State).To(Equal(cluster.StateStable)) + }) + + It("goes ahead past a blocker that can be forced, and records that it was forced", func() { + replica("a", "", "") + workers.set(cluster.WorkerInfo{ID: "old", Name: "old"}) + got, _, err := request(cluster.CarrierTunnel, true) + Expect(err).ToNot(HaveOccurred()) + Expect(got.State).To(Equal(cluster.StatePrepare)) + Expect(got.Force).To(BeTrue()) + }) + + It("does not let force start a change to the carrier that is active", func() { + replica("a", "", "") + _, _, err := request(cluster.CarrierNATS, true) + var blocked *cluster.BlockedError + Expect(errors.As(err, &blocked)).To(BeTrue()) + }) + + It("is busy while another change is under way", func() { + replica("a", "", "") + _, _, err := request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + _, _, err = request(cluster.CarrierTunnel, false) + Expect(err).To(MatchError(cluster.ErrBusy)) + }) + + It("lets one of two concurrent requests win", func() { + replica("a", "", "") + var won, busy atomic.Int32 + start := make(chan struct{}) + var wg sync.WaitGroup + for range 4 { + wg.Go(func() { + <-start + _, _, err := request(cluster.CarrierTunnel, false) + switch { + case err == nil: + won.Add(1) + case errors.Is(err, cluster.ErrBusy): + busy.Add(1) + default: + Fail(err.Error()) + } + }) + } + close(start) + wg.Wait() + Expect(won.Load()).To(Equal(int32(1))) + Expect(busy.Load()).To(Equal(int32(3))) + Expect(row().Epoch).To(Equal(int64(2))) + }) + }) + + Describe("an abort", func() { + It("returns a cluster in prepare to stable on the carrier it had", func() { + replica("a", "", "") + _, _, err := request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + + got, err := sw.Abort(ctx, "admin-2") + Expect(err).ToNot(HaveOccurred()) + Expect(got.State).To(Equal(cluster.StateStable)) + Expect(got.Active).To(Equal(cluster.CarrierNATS)) + Expect(got.Target).To(BeEmpty()) + Expect(got.Epoch).To(Equal(int64(3))) + Expect(got.Note).To(ContainSubstring("aborted by admin-2")) + }) + + It("is refused when nothing is being prepared", func() { + _, err := sw.Abort(ctx, "admin") + Expect(err).To(MatchError(cluster.ErrNotAbortable)) + }) + + It("is refused after the commit, because the target is active", func() { + replica("a", "", "") + _, _, err := request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StateCommit)) + _, err = sw.Abort(ctx, "admin") + Expect(err).To(MatchError(cluster.ErrNotAbortable)) + }) + }) + + Describe("the leader", func() { + prepare := func(force bool) { + GinkgoHelper() + _, _, err := request(cluster.CarrierTunnel, force) + Expect(err).ToNot(HaveOccurred()) + } + + It("does nothing to a stable cluster", func() { + replica("a", "", "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().Epoch).To(Equal(int64(1))) + }) + + It("commits when every live replica is ready for the epoch of the prepare", func() { + replica("a", "", "") + replica("b", "", "") + prepare(false) + epoch := row().Epoch + ready("a", epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StatePrepare), "b is not ready") + + ready("b", epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + got := row() + Expect(got.State).To(Equal(cluster.StateCommit)) + Expect(got.Active).To(Equal(cluster.CarrierTunnel)) + Expect(got.Draining).To(Equal(cluster.CarrierNATS)) + Expect(got.DrainingUntil).ToNot(BeNil()) + Expect(got.Epoch).To(Equal(epoch + 1)) + }) + + It("does not count a replica that is ready for another epoch", func() { + replica("a", "", "") + prepare(false) + ready("a", row().Epoch-1, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StatePrepare)) + }) + + It("ignores a replica that stopped heartbeating", func() { + replica("a", "", "") + replica("b", "", "") + prepare(false) + ready("a", row().Epoch, "") + Expect(db.Exec(`UPDATE instances SET last_seen = now() - interval '10 minutes' WHERE id = 'b'`).Error).To(Succeed()) + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StateCommit)) + }) + + It("aborts at once when a replica says it cannot build the target", func() { + replica("a", "", "") + replica("b", "", "") + prepare(false) + epoch := row().Epoch + ready("a", epoch, "") + ready("b", epoch, "cannot reach the NATS server") + Expect(sw.Drive(ctx)).To(Succeed()) + got := row() + Expect(got.State).To(Equal(cluster.StateStable)) + Expect(got.Active).To(Equal(cluster.CarrierNATS)) + Expect(got.Note).To(ContainSubstring("b")) + Expect(got.Note).To(ContainSubstring("cannot reach the NATS server")) + }) + + It("aborts when a replica is not ready within the prepare timeout, and names it", func() { + timings.PrepareTimeout = 300 * time.Millisecond + replica("a", "", "") + replica("b", "", "") + prepare(false) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StatePrepare)) + + Eventually(func() cluster.State { + Expect(sw.Drive(ctx)).To(Succeed()) + return row().State + }, "10s", "100ms").Should(Equal(cluster.StateStable)) + got := row() + Expect(got.Active).To(Equal(cluster.CarrierNATS)) + Expect(got.Note).To(ContainSubstring("b")) + Expect(got.Note).To(ContainSubstring("not ready")) + }) + + It("commits without a replica that is late, when the request was forced, and names it", func() { + timings.PrepareTimeout = 300 * time.Millisecond + replica("a", "", "") + replica("b", "", "") + prepare(true) + ready("a", row().Epoch, "") + + Eventually(func() cluster.State { + Expect(sw.Drive(ctx)).To(Succeed()) + return row().State + }, "10s", "100ms").Should(Equal(cluster.StateCommit)) + got := row() + Expect(got.Active).To(Equal(cluster.CarrierTunnel)) + Expect(got.Note).To(ContainSubstring("without")) + Expect(got.Note).To(ContainSubstring("b")) + }) + + It("does not wait for the timeout for a forced change when a replica says it cannot build the target", func() { + replica("a", "", "") + replica("b", "", "") + prepare(true) + epoch := row().Epoch + ready("a", epoch, "") + ready("b", epoch, "cannot reach the NATS server") + Expect(sw.Drive(ctx)).To(Succeed()) + got := row() + Expect(got.State).To(Equal(cluster.StateCommit)) + Expect(got.Note).To(ContainSubstring("b")) + }) + + It("settles to stable when every live replica confirms the commit", func() { + replica("a", "", "") + replica("b", "", "") + prepare(false) + ready("a", row().Epoch, "") + ready("b", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + commit := row() + ready("a", commit.Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StateCommit), "b has not confirmed") + + ready("b", commit.Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + got := row() + Expect(got.State).To(Equal(cluster.StateStable)) + Expect(got.Active).To(Equal(cluster.CarrierTunnel)) + Expect(got.Draining).To(Equal(cluster.CarrierNATS), "the previous carrier stays attached") + Expect(got.DrainingUntil).To(Equal(commit.DrainingUntil)) + }) + + It("settles to stable after the transition window when a replica never confirms", func() { + timings.TransitionWindow = 300 * time.Millisecond + replica("a", "", "") + prepare(false) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Eventually(func() cluster.State { + Expect(sw.Drive(ctx)).To(Succeed()) + return row().State + }, "10s", "100ms").Should(Equal(cluster.StateStable)) + Expect(row().Note).To(ContainSubstring("a")) + }) + + It("ends the drain when its time is over, and not before", func() { + timings.MaxDrain = 500 * time.Millisecond + replica("a", "", "") + prepare(false) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + stable := row() + Expect(stable.State).To(Equal(cluster.StateStable)) + Expect(stable.Draining).To(Equal(cluster.CarrierNATS)) + + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().Epoch).To(Equal(stable.Epoch), "the drain is not over") + + Eventually(func() cluster.Carrier { + Expect(sw.Drive(ctx)).To(Succeed()) + return row().Draining + }, "10s", "100ms").Should(BeEmpty()) + got := row() + Expect(got.State).To(Equal(cluster.StateStable)) + Expect(got.DrainingUntil).To(BeNil()) + Expect(got.Active).To(Equal(cluster.CarrierTunnel)) + Expect(got.Epoch).To(Equal(stable.Epoch + 1)) + }) + + It("makes one transition when two leaders run the same step", func() { + replica("a", "", "") + prepare(false) + ready("a", row().Epoch, "") + epoch := row().Epoch + var wg sync.WaitGroup + for range 4 { + wg.Go(func() { _ = sw.Drive(ctx) }) + } + wg.Wait() + Expect(row().Epoch).To(Equal(epoch+1), "the step is a compare-and-set, so a second leader changes nothing") + Expect(row().State).To(Equal(cluster.StateCommit)) + }) + + It("goes on from the row alone, as a leader that takes over would", func() { + replica("a", "", "") + prepare(false) + ready("a", row().Epoch, "") + // A new Switch over the same database has no memory of the first. + other, err := cluster.NewSwitch(cluster.SwitchOptions{Store: store, Registry: reg, Timings: func() cluster.Timings { return timings }}) + Expect(err).ToNot(HaveOccurred()) + Expect(other.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StateCommit)) + }) + + It("keeps the previous drain when a change starts and ends before it has drained", func() { + replica("a", "", "") + prepare(false) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + draining := row() + + _, _, err := request(cluster.CarrierNATS, false) + Expect(err).ToNot(HaveOccurred()) + Expect(row().Draining).To(Equal(draining.Draining), "the previous drain is not forgotten by the next prepare") + _, err = sw.Abort(ctx, "admin") + Expect(err).ToNot(HaveOccurred()) + Expect(row().Draining).To(Equal(draining.Draining)) + Expect(row().DrainingUntil).ToNot(BeNil()) + }) + }) +}) From 6270d6230399000127738d03a9da51b62c6287d0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:29:04 +0000 Subject: [PATCH 107/162] feat(cluster): tell a hook about every move of the carrier row The leader and the admin move the row, and the other replicas read it on a poll. A hook that is called after each move lets the application send the other replicas a hint, so that they look at the row at once instead of at their next poll. The hint is a courtesy, and the poll stays the authority. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/switch.go | 27 +++++++++++++++++++++------ core/services/cluster/switch_test.go | 26 ++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index 83cedab8457f..d5d81482723a 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -198,7 +198,12 @@ type SwitchOptions struct { // AvailabilityMaxAge is how old a replica's report of what it can build may // be. Two minutes when zero. AvailabilityMaxAge time.Duration - Meter metric.Meter + // OnChange is called after every move of the row, with the new row. The + // application uses it to send a hint to the other replicas, so that they look + // at the row at once. It is a courtesy: a replica that never hears it reads + // the row at its next poll. It may be nil. + OnChange func(CarrierRow) + Meter metric.Meter } // Switch is the protocol of a change of carrier. @@ -236,6 +241,16 @@ func NewSwitch(o SwitchOptions) (*Switch, error) { return s, nil } +// transition is CarrierStore.Transition, and tells OnChange about a move that was +// made. +func (s *Switch) transition(ctx context.Context, from int64, change Change) (CarrierRow, error) { + row, err := s.o.Store.Transition(ctx, from, change) + if err == nil && s.o.OnChange != nil { + s.o.OnChange(row) + } + return row, err +} + func (s *Switch) count(outcome string) { if s.changes != nil { s.changes.Add(context.Background(), 1, metric.WithAttributes(attribute.String("outcome", outcome))) @@ -380,7 +395,7 @@ func (s *Switch) Request(ctx context.Context, req Request) (CarrierRow, Report, if req.Force && len(report.Blockers) > 0 { note = fmt.Sprintf("forced past %d blocker(s)", len(report.Blockers)) } - next, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + next, err := s.transition(ctx, row.Epoch, Change{ Active: row.Active, State: StatePrepare, Target: req.Target, Draining: row.Draining, DrainingUntil: row.DrainingUntil, Force: req.Force, Note: note, By: req.By, @@ -414,7 +429,7 @@ func (s *Switch) Abort(ctx context.Context, by string) (CarrierRow, error) { } func (s *Switch) abort(ctx context.Context, row CarrierRow, by, note string) (CarrierRow, error) { - next, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + next, err := s.transition(ctx, row.Epoch, Change{ Active: row.Active, State: StateStable, Draining: row.Draining, DrainingUntil: row.DrainingUntil, Note: note, By: by, @@ -486,7 +501,7 @@ func (s *Switch) drivePrepare(ctx context.Context, row CarrierRow, now time.Time commit := func(note string) error { until := now.Add(t.MaxDrain) - _, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + _, err := s.transition(ctx, row.Epoch, Change{ Active: row.Target, State: StateCommit, Target: row.Target, Draining: row.Active, DrainingUntil: &until, Force: row.Force, Note: note, By: row.ChangedBy, @@ -537,7 +552,7 @@ func (s *Switch) driveCommit(ctx context.Context, row CarrierRow, age time.Durat if len(missing) > 0 { note = strings.TrimSpace(note + " The replica did not confirm the commit within " + t.TransitionWindow.String() + ": " + strings.Join(missing, ", ")) } - _, err = s.o.Store.Transition(ctx, row.Epoch, Change{ + _, err = s.transition(ctx, row.Epoch, Change{ Active: row.Active, State: StateStable, Draining: row.Draining, DrainingUntil: row.DrainingUntil, Force: row.Force, Note: note, By: row.ChangedBy, @@ -557,7 +572,7 @@ func (s *Switch) driveDrain(ctx context.Context, row CarrierRow, now time.Time) if row.DrainingUntil != nil && now.Before(*row.DrainingUntil) { return nil } - _, err := s.o.Store.Transition(ctx, row.Epoch, Change{ + _, err := s.transition(ctx, row.Epoch, Change{ Active: row.Active, State: StateStable, Note: "drained", By: row.ChangedBy, }) if err == nil { diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go index 00adedc892a4..ae393ae62c19 100644 --- a/core/services/cluster/switch_test.go +++ b/core/services/cluster/switch_test.go @@ -309,6 +309,32 @@ var _ = Describe("The switch of the carrier", func() { }) }) + Describe("the hint", func() { + It("tells the hook about every move of the row, and about no move that failed", func() { + var epochs []int64 + var mu sync.Mutex + hinted, err := cluster.NewSwitch(cluster.SwitchOptions{ + Store: store, Registry: reg, Workers: workers, Work: work, + Timings: func() cluster.Timings { return timings }, + OnChange: func(r cluster.CarrierRow) { mu.Lock(); epochs = append(epochs, r.Epoch); mu.Unlock() }, + }) + Expect(err).ToNot(HaveOccurred()) + replica("a", "", "") + + _, _, err = hinted.Request(ctx, cluster.Request{Target: cluster.CarrierTunnel, By: "admin"}) + Expect(err).ToNot(HaveOccurred()) + _, _, err = hinted.Request(ctx, cluster.Request{Target: cluster.CarrierTunnel, By: "admin"}) + Expect(err).To(MatchError(cluster.ErrBusy)) + ready("a", row().Epoch, "") + Expect(hinted.Drive(ctx)).To(Succeed()) + Expect(hinted.Drive(ctx)).To(Succeed()) // nothing is due: a has not confirmed + + mu.Lock() + defer mu.Unlock() + Expect(epochs).To(Equal([]int64{2, 3})) + }) + }) + Describe("an abort", func() { It("returns a cluster in prepare to stable on the carrier it had", func() { replica("a", "", "") From 5f04d5a493dde1e468f82c0ca2f0d412a8487376 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:29:04 +0000 Subject: [PATCH 108/162] feat(carrier): swap the set of a replica behind the holders, and drain the old one Swapper is what one replica does in a change of carrier, and it acts on the row alone. In prepare it builds the target set while the old one serves, attaches every subscription to it, and reports ready for the epoch of the row. In commit it stores the new set, so that publishes, enqueues and new calls go to it, starts the work of the new carrier, and confirms. The old set stays attached through the drain, so a replica that has not flipped yet is still heard, and a call that is running is never moved. When the leader ends the drain the replica stops listening on the old set, stops its work, hands over what it held, closes it, and runs the reconnect hooks once. An abort drops the set that was built for nothing. A replica that polls late, restarts, or joins during a change ends where the others do. Window is the one place with per-worker routing, and it exists only between the commit and the end of the drain. A control request, a file transfer, a dial or a client for one worker goes to the active carrier when the worker is attached to it, else to the previous carrier when the worker is attached to that, else it fails with ErrNoRoute. The choice is read from the state of the worker and never from the error of an earlier attempt. Outside the window every call costs one atomic load. A set can now start work when it becomes active (Start) and hand over what it holds when it is released (Handoff). A run that the end of the drain cuts off is cancelled with messaging.ErrCarrierReleased, so that it is not run again on the carrier that took over. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../carrier/broadcaster_conformance_test.go | 38 + core/services/carrier/holders.go | 270 ++++++- core/services/carrier/set.go | 16 + core/services/carrier/swap.go | 523 ++++++++++++++ core/services/carrier/swap_test.go | 662 ++++++++++++++++++ core/services/carrier/window.go | 136 ++++ core/services/carrier/window_test.go | 153 ++++ core/services/messaging/workqueue.go | 12 +- 8 files changed, 1774 insertions(+), 36 deletions(-) create mode 100644 core/services/carrier/swap.go create mode 100644 core/services/carrier/swap_test.go create mode 100644 core/services/carrier/window.go create mode 100644 core/services/carrier/window_test.go diff --git a/core/services/carrier/broadcaster_conformance_test.go b/core/services/carrier/broadcaster_conformance_test.go index 00cb57e0c34f..ad1e6b8089e2 100644 --- a/core/services/carrier/broadcaster_conformance_test.go +++ b/core/services/carrier/broadcaster_conformance_test.go @@ -1,7 +1,9 @@ package carrier_test import ( + "sync" "sync/atomic" + "time" "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" @@ -38,4 +40,40 @@ var _ = Describe("Broadcaster holder conformance", func() { return messagingtest.Carrier{Bus: h, Peer: h, ServesControlRoots: true} }) }) + + // The flip of a change of carrier is three calls: listen on the next set, store + // it, release the old one. This runs them in a loop while the suite runs, so + // every spec of the suite meets the holder in the middle of a flip. + Describe("while the carrier is swapped under it", func() { + messagingtest.RunBroadcasterConformance(func() messagingtest.Carrier { + var cur atomic.Pointer[carrier.Set] + sets := [2]*carrier.Set{newFakeCarrier(cluster.CarrierNATS, 1).set, newFakeCarrier(cluster.CarrierTunnel, 2).set} + cur.Store(sets[0]) + h := carrier.NewBroadcaster(&cur) + + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Go(func() { + for i := 1; ; i++ { + select { + case <-stop: + return + case <-time.After(time.Millisecond): + } + old, next := sets[(i+1)%2], sets[i%2] + if err := h.Listen(next); err != nil { + Fail(err.Error()) + } + cur.Store(next) + if err := h.Release(old); err != nil { + Fail(err.Error()) + } + } + }) + return messagingtest.Carrier{Bus: h, Peer: h, ServesControlRoots: true, Cleanup: func() { + close(stop) + wg.Wait() + }} + }) + }) }) diff --git a/core/services/carrier/holders.go b/core/services/carrier/holders.go index 2388d9c7c7da..4978de3fec64 100644 --- a/core/services/carrier/holders.go +++ b/core/services/carrier/holders.go @@ -2,6 +2,7 @@ package carrier import ( "context" + "errors" "net" "sync/atomic" "time" @@ -18,70 +19,163 @@ import ( // carrier to one set, and a call that is already running is never moved. The // holders take no lock and allocate nothing of their own. -// Commands forwards the control verbs to the current set. -type Commands struct{ cur *atomic.Pointer[Set] } +// Commands forwards the control verbs to the current set. During a change of +// carrier (see Window) a verb for one worker goes to the set that the worker is +// attached to. +type Commands struct { + cur *atomic.Pointer[Set] + win *Window +} var _ nodes.NodeControl = (*Commands)(nil) // NewCommands returns a holder over cur, which must name a set. func NewCommands(cur *atomic.Pointer[Set]) *Commands { return &Commands{cur: cur} } +// UseWindow makes the holder route by the window. Call it before the holder is +// shared. +func (h *Commands) UseWindow(w *Window) { h.win = w } + +// pick returns the set that a verb for nodeID uses. +func (h *Commands) pick(ctx context.Context, nodeID string) (*Set, error) { + cur := h.cur.Load() + if h.win == nil { + return cur, nil + } + return h.win.route(ctx, cur, nodeID) +} + func (h *Commands) InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { - return h.cur.Load().Commands.InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) } func (h *Commands) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error) { - return h.cur.Load().Commands.UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) } func (h *Commands) DeleteBackend(nodeID, backendName string) (*workerctl.BackendDeleteReply, error) { - return h.cur.Load().Commands.DeleteBackend(nodeID, backendName) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.DeleteBackend(nodeID, backendName) } func (h *Commands) ListBackends(nodeID string) (*workerctl.BackendListReply, error) { - return h.cur.Load().Commands.ListBackends(nodeID) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.ListBackends(nodeID) } func (h *Commands) StopBackend(nodeID, backend string) error { - return h.cur.Load().Commands.StopBackend(nodeID, backend) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return err + } + return s.Commands.StopBackend(nodeID, backend) } func (h *Commands) UnloadModelOnNode(nodeID, modelName string) error { - return h.cur.Load().Commands.UnloadModelOnNode(nodeID, modelName) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return err + } + return s.Commands.UnloadModelOnNode(nodeID, modelName) } -func (h *Commands) PingNode(nodeID string) error { return h.cur.Load().Commands.PingNode(nodeID) } +func (h *Commands) PingNode(nodeID string) error { + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return err + } + return s.Commands.PingNode(nodeID) +} func (h *Commands) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { - return h.cur.Load().Commands.InstallBackendOp(nodeID, backendType, modelID, galleriesJSON, replicaIndex, opID, operationID, deadline, onProgress) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.InstallBackendOp(nodeID, backendType, modelID, galleriesJSON, replicaIndex, opID, operationID, deadline, onProgress) } func (h *Commands) StopLoadOperation(ctx context.Context, nodeID string, req workerctl.ModelStopRequest) (workerctl.ModelStopReply, error) { - return h.cur.Load().Commands.StopLoadOperation(ctx, nodeID, req) + s, err := h.pick(ctx, nodeID) + if err != nil { + return workerctl.ModelStopReply{}, err + } + return s.Commands.StopLoadOperation(ctx, nodeID, req) } func (h *Commands) OperationControl(nodeID string, req workerctl.OperationRequest) (*workerctl.OperationReply, error) { - return h.cur.Load().Commands.OperationControl(nodeID, req) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.OperationControl(nodeID, req) } func (h *Commands) UnloadReplica(nodeID string, replica nodes.NodeModel) error { - return h.cur.Load().Commands.UnloadReplica(nodeID, replica) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return err + } + return s.Commands.UnloadReplica(nodeID, replica) } func (h *Commands) StopModelReplica(ctx context.Context, nodeID string, replica nodes.NodeModel, force bool) (workerctl.ModelStopReply, error) { - return h.cur.Load().Commands.StopModelReplica(ctx, nodeID, replica, force) + s, err := h.pick(ctx, nodeID) + if err != nil { + return workerctl.ModelStopReply{}, err + } + return s.Commands.StopModelReplica(ctx, nodeID, replica, force) } func (h *Commands) ListRunningModels(nodeID string) (*workerctl.ModelsRunningReply, error) { - return h.cur.Load().Commands.ListRunningModels(nodeID) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.ListRunningModels(nodeID) } func (h *Commands) UnloadRemoteModel(modelName string) error { - return h.cur.Load().Commands.UnloadRemoteModel(modelName) + return h.UnloadRemoteModelContext(context.Background(), modelName, false) } +// previous returns the set that is draining, or nil outside a change of carrier. +func (h *Commands) previous(cur *Set) *Set { + if h.win == nil { + return nil + } + if prev := h.win.Previous(); prev != cur { + return prev + } + return nil +} + +// UnloadRemoteModelContext stops a model on every node that holds it. Inside the +// window the active carrier does it first: it removes the rows of the nodes it +// reached, so the nodes that are left are the ones attached to the previous +// carrier, and that one is asked for them. Its answer is the answer, because it +// saw only the nodes that the first one could not do. func (h *Commands) UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error { - return h.cur.Load().Commands.UnloadRemoteModelContext(ctx, modelName, force) + cur := h.cur.Load() + err := cur.Commands.UnloadRemoteModelContext(ctx, modelName, force) + if prev := h.previous(cur); prev != nil && err != nil { + return prev.Commands.UnloadRemoteModelContext(ctx, modelName, force) + } + return err } func (h *Commands) HasRemoteModel(ctx context.Context, modelName string) (bool, error) { @@ -90,52 +184,112 @@ func (h *Commands) HasRemoteModel(ctx context.Context, modelName string) (bool, func (h *Commands) InstallTimeout() time.Duration { return h.cur.Load().Commands.InstallTimeout() } +// DeleteModelFiles asks every node that holds the model to delete its files. +// Inside the window both carriers ask: each reaches the nodes that are attached +// to it, and the verb is idempotent. func (h *Commands) DeleteModelFiles(modelName string) error { - return h.cur.Load().Commands.DeleteModelFiles(modelName) + cur := h.cur.Load() + err := cur.Commands.DeleteModelFiles(modelName) + if prev := h.previous(cur); prev != nil { + err = errors.Join(err, prev.Commands.DeleteModelFiles(modelName)) + } + return err } func (h *Commands) InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { - return h.cur.Load().Commands.InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) + s, err := h.pick(context.Background(), nodeID) + if err != nil { + return nil, err + } + return s.Commands.InstallBackendForce(nodeID, backendType, galleriesJSON, uri, name, alias, replicaIndex, opID, onProgress) } -// Files forwards file staging to the current set. -type Files struct{ cur *atomic.Pointer[Set] } +// Files forwards file staging to the current set. During a change of carrier a +// transfer for one worker goes to the set that the worker is attached to. +type Files struct { + cur *atomic.Pointer[Set] + win *Window +} var _ FileCarrier = (*Files)(nil) // NewFiles returns a holder over cur, which must name a set. func NewFiles(cur *atomic.Pointer[Set]) *Files { return &Files{cur: cur} } +// UseWindow makes the holder route by the window. Call it before the holder is +// shared. +func (h *Files) UseWindow(w *Window) { h.win = w } + +func (h *Files) pick(ctx context.Context, nodeID string) (*Set, error) { + cur := h.cur.Load() + if h.win == nil { + return cur, nil + } + return h.win.route(ctx, cur, nodeID) +} + func (h *Files) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) { - return h.cur.Load().Files.EnsureRemote(ctx, nodeID, localPath, key) + s, err := h.pick(ctx, nodeID) + if err != nil { + return "", err + } + return s.Files.EnsureRemote(ctx, nodeID, localPath, key) } func (h *Files) FetchRemote(ctx context.Context, nodeID, remotePath, localDst string) error { - return h.cur.Load().Files.FetchRemote(ctx, nodeID, remotePath, localDst) + s, err := h.pick(ctx, nodeID) + if err != nil { + return err + } + return s.Files.FetchRemote(ctx, nodeID, remotePath, localDst) } func (h *Files) FetchRemoteByKey(ctx context.Context, nodeID, key, localDst string) error { - return h.cur.Load().Files.FetchRemoteByKey(ctx, nodeID, key, localDst) + s, err := h.pick(ctx, nodeID) + if err != nil { + return err + } + return s.Files.FetchRemoteByKey(ctx, nodeID, key, localDst) } func (h *Files) AllocRemoteTemp(ctx context.Context, nodeID string) (string, error) { - return h.cur.Load().Files.AllocRemoteTemp(ctx, nodeID) + s, err := h.pick(ctx, nodeID) + if err != nil { + return "", err + } + return s.Files.AllocRemoteTemp(ctx, nodeID) } func (h *Files) StageRemoteToStore(ctx context.Context, nodeID, remotePath, key string) error { - return h.cur.Load().Files.StageRemoteToStore(ctx, nodeID, remotePath, key) + s, err := h.pick(ctx, nodeID) + if err != nil { + return err + } + return s.Files.StageRemoteToStore(ctx, nodeID, remotePath, key) } func (h *Files) ReleaseRemote(ctx context.Context, nodeID, key string) error { - return h.cur.Load().Files.ReleaseRemote(ctx, nodeID, key) + s, err := h.pick(ctx, nodeID) + if err != nil { + return err + } + return s.Files.ReleaseRemote(ctx, nodeID, key) } func (h *Files) ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error) { - return h.cur.Load().Files.ListRemoteDir(ctx, nodeID, keyPrefix) + s, err := h.pick(ctx, nodeID) + if err != nil { + return nil, err + } + return s.Files.ListRemoteDir(ctx, nodeID, keyPrefix) } func (h *Files) ReleaseRemoteRequest(ctx context.Context, nodeID, requestID string, keys []string) error { - return h.cur.Load().Files.ReleaseRemoteRequest(ctx, nodeID, requestID, keys) + s, err := h.pick(ctx, nodeID) + if err != nil { + return err + } + return s.Files.ReleaseRemoteRequest(ctx, nodeID, requestID, keys) } // WorkQueue forwards Enqueue to the current set. Producers always use the @@ -154,15 +308,32 @@ func (h *WorkQueue) Enqueue(ctx context.Context, kind messaging.WorkKind, payloa // Clients forwards backend client construction to the current set. A client // built before a swap keeps the connection it has; the code that caches // clients drops them when the carrier changes. -type Clients struct{ cur *atomic.Pointer[Set] } +type Clients struct { + cur *atomic.Pointer[Set] + win *Window +} var _ nodes.BackendClientFactory = (*Clients)(nil) // NewClients returns a holder over cur, which must name a set. func NewClients(cur *atomic.Pointer[Set]) *Clients { return &Clients{cur: cur} } +// UseWindow makes the holder build the client of a worker on the carrier that +// the worker is attached to during a change of carrier. Call it before the +// holder is shared. +func (h *Clients) UseWindow(w *Window) { h.win = w } + func (h *Clients) NewClient(nodeID, address string, parallel bool) grpc.Backend { - return h.cur.Load().Clients.NewClient(nodeID, address, parallel) + cur := h.cur.Load() + if h.win != nil { + // A client is built without a way to fail, so a worker that is attached to + // neither carrier gets a client of the active one, whose first call reports + // the missing route. + if s, err := h.win.route(context.Background(), cur, nodeID); err == nil { + cur = s + } + } + return cur.Clients.NewClient(nodeID, address, parallel) } // NewWorkerDialer returns a nodes.WorkerNetDialerFor that picks the set at the @@ -170,23 +341,52 @@ func (h *Clients) NewClient(nodeID, address string, parallel bool) grpc.Backend // keeps its dial function for as long as it lives, so resolving earlier would // pin it to the carrier it was created on. func NewWorkerDialer(cur *atomic.Pointer[Set]) nodes.WorkerNetDialerFor { + return NewRoutedWorkerDialer(cur, nil) +} + +// NewRoutedWorkerDialer is NewWorkerDialer for a dial that also follows the +// window: during a change of carrier it dials through the carrier that the worker +// is attached to. +func NewRoutedWorkerDialer(cur *atomic.Pointer[Set], win *Window) nodes.WorkerNetDialerFor { return func(nodeID string) func(context.Context, string, string) (net.Conn, error) { return func(ctx context.Context, network, addr string) (net.Conn, error) { - return cur.Load().Dialer(nodeID)(ctx, network, addr) + set := cur.Load() + if win != nil { + var err error + if set, err = win.route(ctx, set, nodeID); err != nil { + return nil, err + } + } + return set.Dialer(nodeID)(ctx, network, addr) } } } // Agents forwards MCP requests to the agent workers of the current set. -type Agents struct{ cur *atomic.Pointer[Set] } +type Agents struct { + cur *atomic.Pointer[Set] + win *Window +} // NewAgents returns a holder over cur, which must name a set. func NewAgents(cur *atomic.Pointer[Set]) *Agents { return &Agents{cur: cur} } +// UseWindow makes the holder send a request to the carrier that has agent +// workers during a change of carrier. Call it before the holder is shared. +func (h *Agents) UseWindow(w *Window) { h.win = w } + +func (h *Agents) pick(ctx context.Context) *Set { + cur := h.cur.Load() + if h.win == nil { + return cur + } + return h.win.routeAgents(ctx, cur) +} + func (h *Agents) ExecuteMCPTool(ctx context.Context, req mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error) { - return h.cur.Load().Agents.ExecuteMCPTool(ctx, req) + return h.pick(ctx).Agents.ExecuteMCPTool(ctx, req) } func (h *Agents) DiscoverMCPTools(ctx context.Context, req mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error) { - return h.cur.Load().Agents.DiscoverMCPTools(ctx, req) + return h.pick(ctx).Agents.DiscoverMCPTools(ctx, req) } diff --git a/core/services/carrier/set.go b/core/services/carrier/set.go index 86639e1b39a9..bccf21ffcebe 100644 --- a/core/services/carrier/set.go +++ b/core/services/carrier/set.go @@ -13,6 +13,7 @@ package carrier import ( + "context" "errors" "fmt" @@ -57,6 +58,21 @@ type Set struct { // be nil for a carrier that caches nothing per node. ForgetNode func(nodeID string) + // Start runs what the carrier needs a frontend replica to do while it is + // active or draining, such as claiming queued work. It is called when the + // carrier becomes the active one and is not called again. ctx ends the work, + // and stop waits until it has finished. A context that is cancelled with + // messaging.ErrCarrierReleased tells the work that its carrier was released + // and that it must not hand what it still holds to another carrier to run + // again. It may be nil for a carrier that needs nothing of the kind. + Start func(ctx context.Context) (stop func(), err error) + + // Handoff moves what the carrier holds for itself onto next, when the carrier + // is released, such as the queued work of a carrier whose queue is a table. + // It runs after the work that Start began has stopped, may run more than + // once, and must be safe to repeat. It may be nil. + Handoff func(ctx context.Context, next *Set) error + // Close releases the carrier's connections. It may be nil. Close func() } diff --git a/core/services/carrier/swap.go b/core/services/carrier/swap.go new file mode 100644 index 000000000000..066d42beebd8 --- /dev/null +++ b/core/services/carrier/swap.go @@ -0,0 +1,523 @@ +package carrier + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "sync" + "sync/atomic" + "time" + + "github.com/mudler/xlog" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" +) + +// Builder builds the Set for a carrier, completely and checked, for the cluster +// row it is given. It is the one place that names the carriers. +type Builder func(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*Set, error) + +// RowSource reads the cluster carrier row and the clock of the database. +// *cluster.CarrierStore is one. +type RowSource interface { + Get(ctx context.Context) (cluster.CarrierRow, error) + DBNow(ctx context.Context) (time.Time, error) +} + +// ReadyReporter records that this replica has built the carrier for an epoch of +// the row, or why it could not. *cluster.Membership is one. +type ReadyReporter interface { + ReportReady(ctx context.Context, epoch int64, reason string) error +} + +// DefaultSwapInterval is how often a replica reads the cluster carrier row. +const DefaultSwapInterval = 2 * time.Second + +// SwapperOptions are the inputs of NewSwapper. +type SwapperOptions struct { + // Cur is the pointer the holders read. It names the set in use. + Cur *atomic.Pointer[Set] + // Bus is the broadcaster holder over Cur. + Bus *Broadcaster + // Window routes the calls to workers during a change. It may be nil. + Window *Window + Rows RowSource + Ready ReadyReporter + Build Builder + // Interval is the poll, DefaultSwapInterval when zero. Each wait is varied by + // up to a fifth so that replicas do not read the row at the same moment. + Interval time.Duration + // Settle is how long after the release of a carrier its handoff runs again, + // for the producers that were slow to flip. Twice the interval when zero. + Settle time.Duration + Meter metric.Meter +} + +// Swapper is what a replica does in a change of carrier. It follows the row: +// +// - prepare: build the target set while the old one serves, attach every +// subscription to it, and report ready for the epoch of the row. +// - commit: store the target set, so that publishes, enqueues and new calls go +// to it; start its work; keep the old set attached; confirm. +// - stable with a drain: nothing. The old set stays attached. +// - stable with no drain: release the old set. Stop its work, hand over what it +// held, close it, and run the reconnect hooks once. +// - stable after an abort: drop the set that was built for nothing. +// +// The row is authoritative. A replica that polls late, restarts, or joins during a +// change ends in the same place as the others, because it acts on what the row +// says and not on what it saw before. The poll is the only thing that must work; +// the hint only shortens the wait. +type Swapper struct { + o SwapperOptions + + mu sync.Mutex + base context.Context + prepared *Set // built for the change in prepare, attached, not yet in use + draining *Set // the previous set, attached through the drain + members map[*Set]*member + reported struct { + set bool + epoch int64 + reason string + } + prepareSeen time.Time + closed bool + + wake chan struct{} + done chan struct{} + doneOnce sync.Once + wg sync.WaitGroup + + row atomic.Pointer[cluster.CarrierRow] + drainRemaining atomic.Int64 + ready atomic.Bool + swaps metric.Int64Counter + failures metric.Int64Counter + swapSeconds metric.Float64Histogram +} + +// member is the work that a set started. +type member struct { + cancel context.CancelCauseFunc + stop func() +} + +// NewSwapper returns the swapper of a replica. Cur must already name the set +// that the row names. +func NewSwapper(o SwapperOptions) (*Swapper, error) { + switch { + case o.Cur == nil || o.Cur.Load() == nil: + return nil, errors.New("a swapper needs the pointer to the set in use, and it must name a set") + case o.Bus == nil: + return nil, errors.New("a swapper needs the broadcaster holder") + case o.Rows == nil || o.Ready == nil || o.Build == nil: + return nil, errors.New("a swapper needs a row source, a readiness reporter and a builder") + } + if o.Interval <= 0 { + o.Interval = DefaultSwapInterval + } + if o.Settle <= 0 { + o.Settle = 2 * o.Interval + } + s := &Swapper{ + o: o, base: context.Background(), + members: map[*Set]*member{}, + wake: make(chan struct{}, 1), + done: make(chan struct{}), + } + s.registerMetrics() + return s, nil +} + +func (s *Swapper) registerMetrics() { + meter := s.o.Meter + if meter == nil { + meter = otel.Meter("github.com/mudler/LocalAI") + } + s.swaps, _ = meter.Int64Counter("localai_carrier_swaps_total", + metric.WithDescription("Times this replica started to use another carrier, by the carrier it moved to")) + s.failures, _ = meter.Int64Counter("localai_carrier_prepare_failures_total", + metric.WithDescription("Times this replica could not build the carrier it was asked to prepare")) + s.swapSeconds, _ = meter.Float64Histogram("localai_carrier_swap_seconds", + metric.WithDescription("Time from the first sight of a change in prepare to the start of use of the new carrier on this replica"), metric.WithUnit("s")) + + epoch, _ := meter.Int64ObservableGauge("localai_carrier_epoch", metric.WithDescription("Epoch of the cluster carrier row as this replica last read it")) + state, _ := meter.Int64ObservableGauge("localai_carrier_state", metric.WithDescription("1 for the state the cluster is in (stable, prepare or commit), 0 for the others")) + active, _ := meter.Int64ObservableGauge("localai_carrier_active", metric.WithDescription("1 for the carrier this replica publishes on, 0 for the other")) + drain, _ := meter.Float64ObservableGauge("localai_carrier_drain_remaining_seconds", metric.WithDescription("Seconds until the previous carrier is released, 0 when none is draining"), metric.WithUnit("s")) + ready, _ := meter.Int64ObservableGauge("localai_carrier_replica_ready", metric.WithDescription("1 when this replica has reported ready for the epoch it last read")) + _, _ = meter.RegisterCallback(func(_ context.Context, o metric.Observer) error { + row := s.row.Load() + if row == nil { + return nil + } + o.ObserveInt64(epoch, row.Epoch) + for _, st := range []cluster.State{cluster.StateStable, cluster.StatePrepare, cluster.StateCommit} { + v := int64(0) + if row.State == st { + v = 1 + } + o.ObserveInt64(state, v, metric.WithAttributes(attribute.String("state", string(st)))) + } + cur := s.o.Cur.Load().Name + for _, c := range []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel} { + v := int64(0) + if cur == c { + v = 1 + } + o.ObserveInt64(active, v, metric.WithAttributes(attribute.String("carrier", string(c)))) + } + o.ObserveFloat64(drain, time.Duration(s.drainRemaining.Load()).Seconds()) + r := int64(0) + if s.ready.Load() { + r = 1 + } + o.ObserveInt64(ready, r) + return nil + }, epoch, state, active, drain, ready) +} + +// Adopt starts the work of the set in use, and takes ctx as the parent of every +// piece of work that the swapper starts. Call it once, when the replica starts. +func (s *Swapper) Adopt(ctx context.Context) { + s.mu.Lock() + defer s.mu.Unlock() + s.base = ctx + s.ensureStarted(s.o.Cur.Load()) +} + +// Wake asks the polling loop to look at the row now. It is the hint that the +// leader sends; losing it costs one interval. +func (s *Swapper) Wake() { + select { + case s.wake <- struct{}{}: + default: + } +} + +// Run polls the row until ctx ends. +func (s *Swapper) Run(ctx context.Context) { + defer s.doneOnce.Do(func() { close(s.done) }) + s.Adopt(ctx) + var lastErr string + for { + if err := s.Poll(ctx); err != nil && ctx.Err() == nil { + // One line when it starts and one when it clears, not one per poll. + if msg := err.Error(); msg != lastErr { + lastErr = msg + xlog.Warn("Could not follow the cluster carrier row", "error", err) + } + } else if err == nil { + lastErr = "" + } + wait := time.Duration(float64(s.o.Interval) * (0.8 + 0.4*rand.Float64())) + timer := time.NewTimer(wait) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-s.wake: + timer.Stop() + case <-timer.C: + } + } +} + +// Poll reads the row and acts on it. +func (s *Swapper) Poll(ctx context.Context) error { + row, err := s.o.Rows.Get(ctx) + if errors.Is(err, cluster.ErrNotSeeded) { + return nil + } + if err != nil { + return err + } + if row.Draining != "" && row.DrainingUntil != nil { + if now, err := s.o.Rows.DBNow(ctx); err == nil { + s.drainRemaining.Store(int64(max(0, row.DrainingUntil.Sub(now)))) + } + } else { + s.drainRemaining.Store(0) + } + return s.Reconcile(ctx, row) +} + +// Reconcile acts on one reading of the row. Calls are serialised. +func (s *Swapper) Reconcile(ctx context.Context, row cluster.CarrierRow) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.closed { + return nil + } + s.row.Store(&row) + if row.State == cluster.StatePrepare && s.prepareSeen.IsZero() { + s.prepareSeen = time.Now() + } + + // Be on the carrier that the row says is active before anything else. A + // replica that polled late, or one that missed a whole change, is behind. + if cur := s.o.Cur.Load(); cur.Name != row.Active { + if err := s.swapTo(ctx, row); err != nil { + s.countFailure() + xlog.Warn("This replica cannot use the active carrier of the cluster", "carrier", row.Active, "epoch", row.Epoch, "error", err) + return s.report(ctx, row.Epoch, fmt.Sprintf("cannot build the %s carrier: %v", row.Active, err)) + } + } + s.ensureStarted(s.o.Cur.Load()) + + switch row.State { + case cluster.StatePrepare: + return s.prepare(ctx, row) + case cluster.StateCommit: + return s.report(ctx, row.Epoch, "") + default: + s.prepareSeen = time.Time{} + s.settle(row) + return nil + } +} + +func (s *Swapper) countFailure() { + if s.failures != nil { + s.failures.Add(context.Background(), 1) + } +} + +// report tells the cluster which epoch this replica is ready for. A report that +// is the same as the last one that went through is not written again. +func (s *Swapper) report(ctx context.Context, epoch int64, reason string) error { + s.ready.Store(reason == "") + if s.reported.set && s.reported.epoch == epoch && s.reported.reason == reason { + return nil + } + if err := s.o.Ready.ReportReady(ctx, epoch, reason); err != nil { + return fmt.Errorf("reporting readiness for epoch %d: %w", epoch, err) + } + s.reported.set, s.reported.epoch, s.reported.reason = true, epoch, reason + return nil +} + +// attach builds a set for target, or reuses one that is attached already, and +// makes every subscription listen on it. +func (s *Swapper) attach(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*Set, error) { + switch { + case s.prepared != nil && s.prepared.Name == target: + return s.prepared, nil + case s.draining != nil && s.draining.Name == target: + return s.draining, nil + } + set, err := s.o.Build(ctx, row, target) + if err != nil { + return nil, err + } + if err := set.Validate(); err != nil { + return nil, err + } + if err := s.o.Bus.Listen(set); err != nil { + closeSet(set) + return nil, err + } + return set, nil +} + +func (s *Swapper) prepare(ctx context.Context, row cluster.CarrierRow) error { + set, err := s.attach(ctx, row, row.Target) + if err != nil { + s.countFailure() + xlog.Warn("This replica cannot build the carrier it was asked to prepare", "target", row.Target, "epoch", row.Epoch, "error", err) + return s.report(ctx, row.Epoch, err.Error()) + } + if s.prepared != set { + s.prepared = set + xlog.Info("The target carrier is built and attached", "target", row.Target, "epoch", row.Epoch) + } + return s.report(ctx, row.Epoch, "") +} + +// swapTo starts to use the carrier the row names. +func (s *Swapper) swapTo(ctx context.Context, row cluster.CarrierRow) error { + next, err := s.attach(ctx, row, row.Active) + if err != nil { + return err + } + // Listening is idempotent. It matters for a set that was built for a change that + // this replica never saw in prepare. + if err := s.o.Bus.Listen(next); err != nil { + return err + } + old := s.o.Cur.Load() + s.o.Cur.Store(next) + + third := s.draining + if next == s.draining { + third = nil + } + if next == s.prepared { + s.prepared = nil + } + s.draining = old + if s.o.Window != nil { + s.o.Window.Open(old) + } + if third != nil && third != old { + // A set from a drain that was never closed. It is not in use and nothing + // routes to it any more. + s.retire(third, next) + } + s.ensureStarted(next) + if s.swaps != nil { + s.swaps.Add(ctx, 1, metric.WithAttributes(attribute.String("to", string(next.Name)))) + } + took := time.Duration(0) + if !s.prepareSeen.IsZero() { + took = time.Since(s.prepareSeen) + if s.swapSeconds != nil { + s.swapSeconds.Record(ctx, took.Seconds()) + } + } + xlog.Info("This replica now uses another carrier; the previous one stays attached while it drains", + "carrier", next.Name, "previous", old.Name, "epoch", row.Epoch, "since_prepare", took.Round(time.Millisecond)) + return nil +} + +// ensureStarted starts the work of a set once. +func (s *Swapper) ensureStarted(set *Set) { + if set == nil || set.Start == nil { + return + } + if _, ok := s.members[set]; ok { + return + } + ctx, cancel := context.WithCancelCause(s.base) + stop, err := set.Start(ctx) + if err != nil { + cancel(err) + // Tried again at the next poll. + xlog.Error("Could not start the work of the carrier", "carrier", set.Name, "error", err) + return + } + s.members[set] = &member{cancel: cancel, stop: stop} +} + +// settle acts on a stable row: it drops a set built for a change that was +// aborted, and releases the previous set once the leader has ended the drain. +func (s *Swapper) settle(row cluster.CarrierRow) { + if s.prepared != nil { + if s.prepared != s.draining { + closeSet(s.dropListening(s.prepared)) + } + s.prepared = nil + } + if row.Draining == "" && s.draining != nil { + old := s.draining + s.draining = nil + if s.o.Window != nil { + s.o.Window.Close() + } + s.retire(old, s.o.Cur.Load()) + } +} + +// dropListening detaches every subscription from set. +func (s *Swapper) dropListening(set *Set) *Set { + if err := s.o.Bus.Release(set); err != nil { + xlog.Warn("Could not detach the subscriptions from a carrier", "carrier", set.Name, "error", err) + } + return set +} + +// retire releases old after next took over: stop listening on it, stop its work, +// hand over what it held, close it, and run the reconnect hooks once. The slow +// part runs in the background, so that the poll goes on. +func (s *Swapper) retire(old, next *Set) { + s.dropListening(old) + m := s.members[old] + delete(s.members, old) + s.wg.Add(1) + go func() { + defer s.wg.Done() + if m != nil { + // A run that is still going loses its carrier here. It is told why, so + // that it is not offered to the new carrier to run again. + m.cancel(messaging.ErrCarrierReleased) + m.stop() + } + s.handoff(old, next) + closeSet(old) + // Messages across the flip were not ordered, so the consumers read their + // tables once. + s.o.Bus.NotifyReconnect() + select { + case <-time.After(s.o.Settle): + case <-s.done: + } + // A producer that was slow to flip can still have put work on the old + // carrier. The handoff is a compare-and-set, so a second run finds nothing + // that the first took. + if cur := s.o.Cur.Load(); cur != old { + s.handoff(old, cur) + } + }() +} + +func (s *Swapper) handoff(from, to *Set) { + if from.Handoff == nil || to == nil { + return + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(s.base), 2*time.Minute) + defer cancel() + if err := from.Handoff(ctx, to); err != nil { + xlog.Warn("Could not hand over what the released carrier held", "from", from.Name, "to", to.Name, "error", err) + } +} + +// Close stops the work of every set, and closes the sets that are not in use. +// The set in use is closed by whoever built it, after this returns. It is safe to +// call more than once. +func (s *Swapper) Close() { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return + } + s.closed = true + members := s.members + s.members = map[*Set]*member{} + var extra []*Set + for _, set := range []*Set{s.prepared, s.draining} { + if set != nil && set != s.o.Cur.Load() { + extra = append(extra, set) + } + } + if len(extra) == 2 && extra[0] == extra[1] { + extra = extra[:1] + } + s.prepared, s.draining = nil, nil + s.mu.Unlock() + if s.o.Window != nil { + s.o.Window.Close() + } + + for _, m := range members { + m.cancel(context.Canceled) + m.stop() + } + for _, set := range extra { + s.dropListening(set) + closeSet(set) + } + s.doneOnce.Do(func() { close(s.done) }) + s.wg.Wait() +} + +func closeSet(set *Set) { + if set != nil && set.Close != nil { + set.Close() + } +} diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go new file mode 100644 index 000000000000..80910f3a3d4c --- /dev/null +++ b/core/services/carrier/swap_test.go @@ -0,0 +1,662 @@ +package carrier_test + +import ( + "context" + "errors" + "fmt" + "slices" + "sync" + "sync/atomic" + "time" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// swapSubject is a broadcast subject that every carrier serves. +const swapSubject = "state.swap-spec" + +// fakeNet is the network that the replicas of a spec share. There is one bus for +// each carrier, so a message that one replica publishes on a carrier reaches the +// replicas that listen on it. +type fakeNet struct { + mu sync.Mutex + bus map[cluster.Carrier]*testutil.FakeBus +} + +func newFakeNet() *fakeNet { + return &fakeNet{bus: map[cluster.Carrier]*testutil.FakeBus{ + cluster.CarrierNATS: testutil.NewFakeBus(), + cluster.CarrierTunnel: testutil.NewFakeBus(), + }} +} + +// life records what the replica did to the sets it built. +type life struct { + recorder +} + +// replicaNode is one frontend replica of a spec: its own pointer, holders, +// swapper and the sets it built. +type replicaNode struct { + id string + net *fakeNet + cur atomic.Pointer[carrier.Set] + bus *carrier.Broadcaster + window *carrier.Window + cmds *carrier.Commands + queue *carrier.WorkQueue + swapper *carrier.Swapper + life *life + + mu sync.Mutex + builds map[cluster.Carrier]int + buildErr map[cluster.Carrier]error + sets map[cluster.Carrier][]*fakeCarrier + handedTo []cluster.Carrier + attached map[string]carrier.Attachment + reconnect atomic.Int32 +} + +func (n *replicaNode) build(_ context.Context, row cluster.CarrierRow, target cluster.Carrier) (*carrier.Set, error) { + n.mu.Lock() + defer n.mu.Unlock() + if err := n.buildErr[target]; err != nil { + return nil, err + } + n.builds[target]++ + return n.newSet(target, row.Epoch), nil +} + +// newSet builds a fake set over the shared bus of the carrier. The caller holds +// n.mu. +func (n *replicaNode) newSet(name cluster.Carrier, epoch int64) *carrier.Set { + f := newFakeCarrier(name, epoch) + bus := n.net.bus[name] + f.set.Broadcaster = bus + f.set.OnReconnect = bus.OnReconnect + f.set.Start = func(ctx context.Context) (func(), error) { + n.life.rec("start:" + string(name)) + return func() { + released := errors.Is(context.Cause(ctx), messaging.ErrCarrierReleased) + n.life.rec(fmt.Sprintf("stop:%s:released=%t", name, released)) + }, nil + } + f.set.Handoff = func(_ context.Context, next *carrier.Set) error { + n.mu.Lock() + n.handedTo = append(n.handedTo, next.Name) + n.mu.Unlock() + n.life.rec("handoff:" + string(name) + "->" + string(next.Name)) + return nil + } + f.set.Close = func() { n.life.rec("close:" + string(name)) } + n.sets[name] = append(n.sets[name], f) + return f.set +} + +func (n *replicaNode) fake(name cluster.Carrier) *fakeCarrier { + n.mu.Lock() + defer n.mu.Unlock() + if len(n.sets[name]) == 0 { + return nil + } + return n.sets[name][len(n.sets[name])-1] +} + +func (n *replicaNode) buildCount(name cluster.Carrier) int { + n.mu.Lock() + defer n.mu.Unlock() + return n.builds[name] +} + +func (n *replicaNode) failBuild(name cluster.Carrier, err error) { + n.mu.Lock() + defer n.mu.Unlock() + n.buildErr[name] = err +} + +func (n *replicaNode) active() cluster.Carrier { return n.cur.Load().Name } + +func (n *replicaNode) handoffs() []cluster.Carrier { + n.mu.Lock() + defer n.mu.Unlock() + return slices.Clone(n.handedTo) +} + +var _ = Describe("The swap of a replica", func() { + var ( + ctx context.Context + db *gorm.DB + store *cluster.CarrierStore + reg *cluster.Registry + sw *cluster.Switch + net *fakeNet + nodes []*replicaNode + timings cluster.Timings + ) + + // join starts a replica on the carrier that the row names, as initDistributed does. + join := func(id string) *replicaNode { + GinkgoHelper() + n := &replicaNode{ + id: id, net: net, life: &life{}, + builds: map[cluster.Carrier]int{}, buildErr: map[cluster.Carrier]error{}, + sets: map[cluster.Carrier][]*fakeCarrier{}, attached: map[string]carrier.Attachment{}, + } + row, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(reg.Register(ctx, id, "v1", 0, "")).To(Succeed()) + n.mu.Lock() + n.cur.Store(n.newSet(row.Active, row.Epoch)) + n.mu.Unlock() + n.bus = carrier.NewBroadcaster(&n.cur) + n.window = carrier.NewWindow(func(_ context.Context, nodeID string) (carrier.Attachment, error) { + n.mu.Lock() + defer n.mu.Unlock() + return n.attached[nodeID], nil + }, nil) + n.cmds = carrier.NewCommands(&n.cur) + n.cmds.UseWindow(n.window) + n.queue = carrier.NewWorkQueue(&n.cur) + n.swapper, err = carrier.NewSwapper(carrier.SwapperOptions{ + Cur: &n.cur, Bus: n.bus, Window: n.window, Rows: store, + Ready: readyReporter{reg: reg, id: id}, Build: n.build, + Interval: time.Hour, Settle: 50 * time.Millisecond, + }) + Expect(err).ToNot(HaveOccurred()) + n.bus.OnReconnect(func() { n.reconnect.Add(1) }) + n.swapper.Adopt(ctx) + nodes = append(nodes, n) + return n + } + + row := func() cluster.CarrierRow { + GinkgoHelper() + r, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + return r + } + // poll is one pass of every replica over the row. + poll := func(which ...*replicaNode) { + GinkgoHelper() + if len(which) == 0 { + which = nodes + } + for _, n := range which { + Expect(n.swapper.Poll(ctx)).To(Succeed(), n.id) + } + } + drive := func() { GinkgoHelper(); Expect(sw.Drive(ctx)).To(Succeed()) } + // availability makes every live replica report that it can use both carriers. + availability := func() { + GinkgoHelper() + for _, n := range nodes { + Expect(reg.ReportAvailability(ctx, n.id, map[cluster.Carrier]string{cluster.CarrierNATS: "", cluster.CarrierTunnel: ""})).To(Succeed()) + } + } + request := func(target cluster.Carrier) { + GinkgoHelper() + availability() + _, _, err := sw.Request(ctx, cluster.Request{Target: target, By: "admin"}) + Expect(err).ToNot(HaveOccurred()) + } + // settle runs a whole change with every replica prompt. + settle := func(target cluster.Carrier) { + GinkgoHelper() + request(target) + poll() + drive() // commit + poll() + drive() // stable, draining + Expect(row().State).To(Equal(cluster.StateStable)) + Expect(row().Active).To(Equal(target)) + } + expire := func() { + GinkgoHelper() + Expect(db.Exec(`UPDATE cluster_carrier SET draining_until = now() - interval '1 second'`).Error).To(Succeed()) + drive() + Expect(row().Draining).To(BeEmpty()) + poll() + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + var err error + store, err = cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(ctx, cluster.CarrierNATS, "seed") + Expect(err).ToNot(HaveOccurred()) + reg = cluster.NewRegistry(db) + net = newFakeNet() + nodes = nil + timings = cluster.Timings{PrepareTimeout: time.Minute, TransitionWindow: time.Minute, MaxDrain: time.Minute} + sw, err = cluster.NewSwitch(cluster.SwitchOptions{Store: store, Registry: reg, Timings: func() cluster.Timings { return timings }}) + Expect(err).ToNot(HaveOccurred()) + }) + + AfterEach(func() { + for _, n := range nodes { + n.swapper.Close() + } + }) + + Describe("prepare", func() { + It("builds the target while the old carrier still serves, listens on both, and reports ready", func() { + a, b := join("a"), join("b") + var heardOnNew atomic.Int32 + _, err := a.bus.Subscribe(swapSubject, func([]byte) { heardOnNew.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + + request(cluster.CarrierTunnel) + poll() + + Expect(a.buildCount(cluster.CarrierTunnel)).To(Equal(1)) + Expect(a.active()).To(Equal(cluster.CarrierNATS), "publishing is still on the old carrier") + in, err := reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyEpoch).To(Equal(row().Epoch)) + Expect(in.ReadyReason).To(BeEmpty()) + + // A replica that has not swapped yet publishes on the new carrier, and the + // subscription of this one hears it. + Expect(net.bus[cluster.CarrierTunnel].Publish(swapSubject, "x")).To(Succeed()) + Expect(heardOnNew.Load()).To(Equal(int32(1))) + _ = b + }) + + It("builds the target once, however many times it polls", func() { + a := join("a") + request(cluster.CarrierTunnel) + poll() + poll() + poll() + Expect(a.buildCount(cluster.CarrierTunnel)).To(Equal(1)) + }) + + It("reports why it cannot build the target, and the leader aborts naming it", func() { + a, b := join("a"), join("b") + b.failBuild(cluster.CarrierTunnel, errors.New("the LISTEN connection was refused")) + request(cluster.CarrierTunnel) + poll() + + in, err := reg.Get(ctx, "b") + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyReason).To(ContainSubstring("the LISTEN connection was refused")) + + drive() + got := row() + Expect(got.State).To(Equal(cluster.StateStable)) + Expect(got.Active).To(Equal(cluster.CarrierNATS)) + Expect(got.Note).To(ContainSubstring("b")) + poll() + Expect(a.active()).To(Equal(cluster.CarrierNATS)) + Expect(b.active()).To(Equal(cluster.CarrierNATS)) + }) + + It("tries again to build the target after a failure", func() { + b := join("b") + b.failBuild(cluster.CarrierTunnel, errors.New("not yet")) + request(cluster.CarrierTunnel) + poll() + b.failBuild(cluster.CarrierTunnel, nil) + poll() + in, err := reg.Get(ctx, "b") + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyReason).To(BeEmpty()) + Expect(in.ReadyEpoch).To(Equal(row().Epoch)) + }) + + It("drops the prepared set when the change is aborted, and keeps the old carrier", func() { + a := join("a") + var heard atomic.Int32 + _, err := a.bus.Subscribe(swapSubject, func([]byte) { heard.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + request(cluster.CarrierTunnel) + poll() + prepared := a.fake(cluster.CarrierTunnel) + + _, err = sw.Abort(ctx, "admin") + Expect(err).ToNot(HaveOccurred()) + poll() + + Expect(a.life.seen()).To(ContainElement("close:tunnel")) + Expect(a.active()).To(Equal(cluster.CarrierNATS)) + Expect(net.bus[cluster.CarrierTunnel].Publish(swapSubject, "x")).To(Succeed()) + Expect(heard.Load()).To(BeZero(), "the subscription is not attached to the prepared set any more") + Expect(net.bus[cluster.CarrierNATS].Publish(swapSubject, "x")).To(Succeed()) + Expect(heard.Load()).To(Equal(int32(1))) + Expect(prepared).ToNot(BeNil()) + }) + }) + + Describe("commit", func() { + It("stores the new set, so that publishes go to it, and keeps the old carrier attached", func() { + a, b := join("a"), join("b") + var heardOld, heardNew atomic.Int32 + _, err := b.bus.Subscribe(swapSubject, func([]byte) {}) + Expect(err).ToNot(HaveOccurred()) + net.bus[cluster.CarrierNATS].Subscribe(swapSubject, func([]byte) { heardOld.Add(1) }) + net.bus[cluster.CarrierTunnel].Subscribe(swapSubject, func([]byte) { heardNew.Add(1) }) + + request(cluster.CarrierTunnel) + poll() + drive() + Expect(row().State).To(Equal(cluster.StateCommit)) + poll(a) + Expect(a.active()).To(Equal(cluster.CarrierTunnel)) + Expect(b.active()).To(Equal(cluster.CarrierNATS), "b has not polled yet") + + Expect(a.bus.Publish(swapSubject, "from-a")).To(Succeed()) + Expect(heardNew.Load()).To(Equal(int32(1))) + Expect(heardOld.Load()).To(BeZero(), "a publish goes to one carrier only") + // b has not swapped, publishes on the old carrier, and is heard by a, which + // still listens there. + var heardByA atomic.Int32 + _, err = a.bus.Subscribe(swapSubject, func([]byte) { heardByA.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + Expect(b.bus.Publish(swapSubject, "from-b")).To(Succeed()) + Expect(heardByA.Load()).To(Equal(int32(1))) + }) + + It("confirms the commit so the leader can settle", func() { + a, b := join("a"), join("b") + request(cluster.CarrierTunnel) + poll() + drive() + commit := row() + poll() + for _, n := range []*replicaNode{a, b} { + in, err := reg.Get(ctx, n.id) + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyEpoch).To(Equal(commit.Epoch), n.id) + } + drive() + Expect(row().State).To(Equal(cluster.StateStable)) + Expect(row().Draining).To(Equal(cluster.CarrierNATS)) + }) + + It("delivers every message once to a subscriber while the replicas flip one after the other", func() { + pub, sub, third := join("pub"), join("sub"), join("third") + var mu sync.Mutex + got := map[string]int{} + _, err := sub.bus.Subscribe(swapSubject, func(raw []byte) { + mu.Lock() + got[string(raw)]++ + mu.Unlock() + }) + Expect(err).ToNot(HaveOccurred()) + sent := 0 + publish := func(from *replicaNode) { + sent++ + Expect(from.bus.Publish(swapSubject, fmt.Sprintf("m%d", sent))).To(Succeed()) + } + + publish(pub) + request(cluster.CarrierTunnel) + poll() + publish(pub) + drive() + publish(pub) + poll(pub) // pub flips first + publish(pub) + publish(sub) // sub still on the old carrier + poll(sub) + publish(third) + poll(third) + publish(third) + drive() + publish(pub) + expire() + publish(sub) + + mu.Lock() + defer mu.Unlock() + Expect(got).To(HaveLen(sent), "no message is lost") + for m, n := range got { + Expect(n).To(Equal(1), "message %s was delivered %d times", m, n) + } + }) + + It("builds the target at commit for a replica that missed the prepare", func() { + a := join("a") + request(cluster.CarrierTunnel) + poll(a) + late := join("late") // joins during prepare, on the old carrier + drive() // late was not ready, so this is a no-op + Expect(row().State).To(Equal(cluster.StatePrepare)) + poll(late) + drive() + Expect(row().State).To(Equal(cluster.StateCommit)) + Expect(late.buildCount(cluster.CarrierTunnel)).To(Equal(1)) + }) + + It("goes straight to the new carrier for a replica that sees the commit and never saw the prepare", func() { + a := join("a") + request(cluster.CarrierTunnel) + poll(a) + drive() + commit := row() + Expect(commit.State).To(Equal(cluster.StateCommit)) + late := join("late") // starts on the carrier that the row names + Expect(late.active()).To(Equal(cluster.CarrierTunnel)) + poll(late) + Expect(late.buildCount(cluster.CarrierNATS)).To(BeZero(), "it never builds the old carrier") + in, err := reg.Get(ctx, "late") + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyEpoch).To(Equal(commit.Epoch)) + }) + + It("catches up when it polled too late to see the commit", func() { + a, b := join("a"), join("b") + request(cluster.CarrierTunnel) + poll() + drive() + poll(a) + drive() // settle is not due: b has not confirmed + Expect(row().State).To(Equal(cluster.StateCommit)) + Expect(db.Exec(`UPDATE cluster_carrier SET state = 'stable', epoch = epoch + 1`).Error).To(Succeed()) + poll(b) + Expect(b.active()).To(Equal(cluster.CarrierTunnel)) + }) + + It("starts the work of the new carrier when it commits, and keeps the old one running through the drain", func() { + a := join("a") + request(cluster.CarrierTunnel) + poll() + Expect(a.life.seen()).To(ConsistOf("start:nats"), "only the carrier in use runs its work") + drive() + poll() + Expect(a.life.seen()).To(ConsistOf("start:nats", "start:tunnel")) + }) + }) + + Describe("the drain", func() { + It("keeps the old carrier attached until the leader ends the drain", func() { + a := join("a") + var heard atomic.Int32 + _, err := a.bus.Subscribe(swapSubject, func([]byte) { heard.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + settle(cluster.CarrierTunnel) + poll() + + Expect(net.bus[cluster.CarrierNATS].Publish(swapSubject, "late")).To(Succeed()) + Expect(heard.Load()).To(Equal(int32(1)), "a late publisher on the old carrier is still heard") + Expect(a.life.seen()).ToNot(ContainElement(ContainSubstring("close:nats"))) + }) + + It("releases the old carrier when the drain ends: unsubscribes, stops its work, hands over, closes, and replays the hooks once", func() { + a := join("a") + var heard atomic.Int32 + _, err := a.bus.Subscribe(swapSubject, func([]byte) { heard.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + settle(cluster.CarrierTunnel) + expire() + + Eventually(a.life.seen, "5s").Should(ContainElements( + "stop:nats:released=true", "handoff:nats->tunnel", "close:nats")) + Expect(net.bus[cluster.CarrierNATS].Publish(swapSubject, "after")).To(Succeed()) + Expect(heard.Load()).To(BeZero(), "the old carrier is no longer listened on") + Expect(a.reconnect.Load()).To(Equal(int32(1)), "the consumers read their tables once") + + seen := a.life.seen() + Expect(slices.Index(seen, "stop:nats:released=true")).To(BeNumerically("<", slices.Index(seen, "handoff:nats->tunnel")), + "the work is stopped before what it held is handed over") + Expect(slices.Index(seen, "handoff:nats->tunnel")).To(BeNumerically("<", slices.Index(seen, "close:nats"))) + }) + + It("hands over a second time after the producers that lag have had time to flip", func() { + a := join("a") + settle(cluster.CarrierTunnel) + expire() + Eventually(func() int { + n := 0 + for _, s := range a.life.seen() { + if s == "handoff:nats->tunnel" { + n++ + } + } + return n + }, "5s", "20ms").Should(Equal(2)) + }) + + It("reuses the draining set when the admin switches back during the drain", func() { + a := join("a") + settle(cluster.CarrierTunnel) + natsBuilds := a.buildCount(cluster.CarrierNATS) + Expect(natsBuilds).To(Equal(0), "the replica started on NATS and never built it") + + request(cluster.CarrierNATS) + poll() + Expect(a.buildCount(cluster.CarrierNATS)).To(Equal(0), "the old set is still attached, so there is nothing to build") + drive() + poll() + Expect(a.active()).To(Equal(cluster.CarrierNATS)) + Expect(a.life.seen()).ToNot(ContainElement("close:nats")) + }) + + It("closes both sets of a replica that shuts down", func() { + a := join("a") + settle(cluster.CarrierTunnel) + a.swapper.Close() + Eventually(a.life.seen, "5s").Should(ContainElement("close:nats")) + }) + }) + + Describe("the window", func() { + It("sends a control request to the carrier that the worker is attached to", func() { + a := join("a") + a.attached["on-nats"] = carrier.Attachment{NATS: true} + a.attached["on-tunnel"] = carrier.Attachment{Tunnel: true} + a.attached["on-both"] = carrier.Attachment{NATS: true, Tunnel: true} + settle(cluster.CarrierTunnel) + + Expect(a.cmds.PingNode("on-nats")).To(Succeed()) + Expect(a.cmds.PingNode("on-tunnel")).To(Succeed()) + Expect(a.cmds.PingNode("on-both")).To(Succeed()) + Expect(a.fake(cluster.CarrierNATS).commands.seen()).To(HaveLen(1), "only the worker that stayed on NATS") + Expect(a.fake(cluster.CarrierTunnel).commands.seen()).To(HaveLen(2), "the worker on the tunnel and the one on both") + }) + + It("reports a worker that is attached to neither carrier as ErrNoRoute, and calls nothing", func() { + a := join("a") + settle(cluster.CarrierTunnel) + err := a.cmds.PingNode("nowhere") + Expect(err).To(MatchError(ContainSubstring("no route"))) + Expect(a.fake(cluster.CarrierNATS).commands.seen()).To(BeEmpty()) + Expect(a.fake(cluster.CarrierTunnel).commands.seen()).To(BeEmpty()) + }) + + It("uses the active carrier only once the drain is over", func() { + a := join("a") + a.attached["on-nats"] = carrier.Attachment{NATS: true} + settle(cluster.CarrierTunnel) + expire() + Eventually(a.life.seen, "5s").Should(ContainElement("close:nats")) + + Expect(a.cmds.PingNode("on-nats")).To(Succeed()) + Expect(a.fake(cluster.CarrierNATS).commands.seen()).To(BeEmpty(), "the closed carrier gets no call") + Expect(a.fake(cluster.CarrierTunnel).commands.seen()).To(Equal([]string{"PingNode"}), "the active carrier answers for every worker, and its own error says that the worker is not there") + }) + + It("unloads a model on the nodes that each carrier reaches, and reports only what neither could do", func() { + a := join("a") + settle(cluster.CarrierTunnel) + Expect(a.cmds.UnloadRemoteModel("m")).To(Succeed()) + Expect(a.fake(cluster.CarrierTunnel).commands.seen()).To(ContainElement("UnloadRemoteModelContext")) + Expect(a.fake(cluster.CarrierNATS).commands.seen()).To(BeEmpty(), "the active carrier did it all") + }) + }) + + Describe("a call that runs while the set changes", func() { + It("never blocks and never reaches a closed set", func() { + a, b := join("a"), join("b") + stop := make(chan struct{}) + var wg sync.WaitGroup + var calls atomic.Int64 + for range 4 { + wg.Go(func() { + for { + select { + case <-stop: + return + default: + } + Expect(a.bus.Publish(swapSubject, "x")).To(Succeed()) + Expect(a.queue.Enqueue(ctx, messaging.WorkMCPCI, "p")).To(Succeed()) + calls.Add(1) + } + }) + } + settle(cluster.CarrierTunnel) + settle(cluster.CarrierNATS) + close(stop) + wg.Wait() + Expect(calls.Load()).To(BeNumerically(">", 0)) + _ = b + }) + }) + + Describe("the polling loop", func() { + It("follows the row on its own, and wakes at once on a hint", func() { + a := join("a") + loopCtx, cancel := context.WithCancel(ctx) + defer cancel() + fast, err := carrier.NewSwapper(carrier.SwapperOptions{ + Cur: &a.cur, Bus: a.bus, Window: a.window, Rows: store, + Ready: readyReporter{reg: reg, id: "a"}, Build: a.build, + Interval: time.Hour, Settle: 50 * time.Millisecond, + }) + Expect(err).ToNot(HaveOccurred()) + done := make(chan struct{}) + go func() { fast.Run(loopCtx); close(done) }() + DeferCleanup(func() { cancel(); <-done }) + + request(cluster.CarrierTunnel) + // The interval is an hour, so only the hint can make it look. + Eventually(func() int64 { + fast.Wake() + in, err := reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + return in.ReadyEpoch + }, "10s", "50ms").Should(Equal(row().Epoch)) + }) + }) +}) + +// readyReporter writes the readiness of one replica to the instances table, as +// cluster.Membership does. +type readyReporter struct { + reg *cluster.Registry + id string +} + +func (r readyReporter) ReportReady(ctx context.Context, epoch int64, reason string) error { + return r.reg.ReportReady(ctx, r.id, epoch, reason) +} diff --git a/core/services/carrier/window.go b/core/services/carrier/window.go new file mode 100644 index 000000000000..9dcfe03b6dcd --- /dev/null +++ b/core/services/carrier/window.go @@ -0,0 +1,136 @@ +package carrier + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" +) + +// Attachment says which carriers a worker is connected to. +type Attachment struct{ NATS, Tunnel bool } + +// To reports whether the worker is attached to carrier c. +func (a Attachment) To(c cluster.Carrier) bool { + switch c { + case cluster.CarrierNATS: + return a.NATS + case cluster.CarrierTunnel: + return a.Tunnel + } + return false +} + +// AttachmentSource says which carriers a worker is connected to now. +type AttachmentSource func(ctx context.Context, nodeID string) (Attachment, error) + +// AgentSource says whether any agent worker is attached to carrier c. +type AgentSource func(ctx context.Context, c cluster.Carrier) (bool, error) + +// attachmentTTL is how long an answer of the AttachmentSource is reused. The +// window lasts minutes, and the routing of a request must not cost a database +// read when a burst of requests goes to one worker. +const attachmentTTL = time.Second + +// Window is the routing rule of the time in which two carriers are attached: from +// the commit of a change until the end of its drain. Outside it, every call goes +// to the active set and the window costs one atomic load. +// +// Inside it, a call to a worker goes to the active carrier when the worker is +// attached to it, else to the previous carrier when the worker is attached to +// that one, else nowhere (nodes.ErrNoRoute). The choice is read from the state of +// the worker. It is never made from the error of an earlier attempt, and a call +// is never tried on the other carrier. +type Window struct { + prev atomic.Pointer[Set] + attached AttachmentSource + agents AgentSource + + mu sync.Mutex + cache map[string]cachedAttachment +} + +type cachedAttachment struct { + at time.Time + a Attachment +} + +// NewWindow returns a closed window. agents may be nil: the agent control then +// always uses the active set. +func NewWindow(attached AttachmentSource, agents AgentSource) *Window { + return &Window{attached: attached, agents: agents, cache: map[string]cachedAttachment{}} +} + +// Open starts the window with prev as the carrier that is draining. +func (w *Window) Open(prev *Set) { + w.mu.Lock() + w.cache = map[string]cachedAttachment{} + w.mu.Unlock() + w.prev.Store(prev) +} + +// Close ends the window. +func (w *Window) Close() { w.prev.Store(nil) } + +// Previous returns the set that is draining, or nil outside the window. +func (w *Window) Previous() *Set { return w.prev.Load() } + +func (w *Window) attachment(ctx context.Context, nodeID string) (Attachment, error) { + w.mu.Lock() + hit, ok := w.cache[nodeID] + w.mu.Unlock() + if ok && time.Since(hit.at) < attachmentTTL { + return hit.a, nil + } + a, err := w.attached(ctx, nodeID) + if err != nil { + return Attachment{}, err + } + w.mu.Lock() + w.cache[nodeID] = cachedAttachment{at: time.Now(), a: a} + w.mu.Unlock() + return a, nil +} + +// route returns the set that a call to nodeID uses. +func (w *Window) route(ctx context.Context, cur *Set, nodeID string) (*Set, error) { + prev := w.prev.Load() + if prev == nil || prev == cur || w.attached == nil { + return cur, nil + } + a, err := w.attachment(ctx, nodeID) + if err != nil { + // Not knowing where a worker is attached says nothing about the worker. + // The active carrier is the best guess, and its own error is the answer. + return cur, nil + } + switch { + case a.To(cur.Name): + return cur, nil + case a.To(prev.Name): + return prev, nil + } + return nil, fmt.Errorf("%w: worker %s is attached to neither the %s carrier nor the %s carrier that is draining", + nodes.ErrNoRoute, nodeID, cur.Name, prev.Name) +} + +// routeAgents returns the set that a request for an agent worker uses: the active +// one unless no agent worker is attached to it and one is attached to the +// previous one. +func (w *Window) routeAgents(ctx context.Context, cur *Set) *Set { + prev := w.prev.Load() + if prev == nil || prev == cur || w.agents == nil { + return cur + } + if on, err := w.agents(ctx, cur.Name); err != nil || on { + return cur + } + if on, err := w.agents(ctx, prev.Name); err == nil && on { + return prev + } + return cur +} diff --git a/core/services/carrier/window_test.go b/core/services/carrier/window_test.go new file mode 100644 index 000000000000..898618ad6da4 --- /dev/null +++ b/core/services/carrier/window_test.go @@ -0,0 +1,153 @@ +package carrier_test + +import ( + "context" + "errors" + "sync/atomic" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/nodes" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("The window of a change of carrier", func() { + var ( + old, next *fakeCarrier + cur atomic.Pointer[carrier.Set] + attached map[string]carrier.Attachment + agentsOn map[cluster.Carrier]bool + window *carrier.Window + ) + ctx := context.Background() + + BeforeEach(func() { + old = newFakeCarrier(cluster.CarrierNATS, 1) + next = newFakeCarrier(cluster.CarrierTunnel, 2) + cur = atomic.Pointer[carrier.Set]{} + cur.Store(next.set) + attached = map[string]carrier.Attachment{} + agentsOn = map[cluster.Carrier]bool{} + window = carrier.NewWindow( + func(_ context.Context, nodeID string) (carrier.Attachment, error) { + if nodeID == "unknown" { + return carrier.Attachment{}, errors.New("the registry is down") + } + return attached[nodeID], nil + }, + func(_ context.Context, c cluster.Carrier) (bool, error) { return agentsOn[c], nil }, + ) + }) + + Describe("when it is closed", func() { + It("sends every call to the active set, without asking where the worker is", func() { + asked := false + closed := carrier.NewWindow(func(context.Context, string) (carrier.Attachment, error) { + asked = true + return carrier.Attachment{}, nil + }, nil) + files := carrier.NewFiles(&cur) + files.UseWindow(closed) + _, err := files.EnsureRemote(ctx, "n", "p", "k") + Expect(err).ToNot(HaveOccurred()) + Expect(next.files.seen()).To(Equal([]string{"EnsureRemote"})) + Expect(asked).To(BeFalse()) + }) + }) + + Describe("when it is open", func() { + BeforeEach(func() { window.Open(old.set) }) + + It("sends file transfers to the carrier of the worker", func() { + attached["a"] = carrier.Attachment{NATS: true} + attached["b"] = carrier.Attachment{Tunnel: true} + files := carrier.NewFiles(&cur) + files.UseWindow(window) + _, err := files.EnsureRemote(ctx, "a", "p", "k") + Expect(err).ToNot(HaveOccurred()) + _, err = files.EnsureRemote(ctx, "b", "p", "k") + Expect(err).ToNot(HaveOccurred()) + Expect(old.files.seen()).To(HaveLen(1)) + Expect(next.files.seen()).To(HaveLen(1)) + + _, err = files.EnsureRemote(ctx, "nowhere", "p", "k") + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeTrue()) + }) + + It("builds the client of a worker on the carrier of the worker", func() { + attached["a"] = carrier.Attachment{NATS: true} + clients := carrier.NewClients(&cur) + clients.UseWindow(window) + clients.NewClient("a", "addr", false) + clients.NewClient("b", "addr", false) + Expect(old.clients.seen()).To(HaveLen(1)) + Expect(next.clients.seen()).To(HaveLen(1), "a worker that is attached to neither gets the active carrier, whose first call says there is no route") + }) + + It("dials a worker through the carrier of the worker, and refuses one that is on neither", func() { + attached["a"] = carrier.Attachment{NATS: true} + dial := carrier.NewRoutedWorkerDialer(&cur, window) + _, _ = dial("a")(ctx, "tcp", "x:1") + Expect(old.dialed.seen()).To(Equal([]string{"dial:a"})) + _, err := dial("nowhere")(ctx, "tcp", "x:1") + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeTrue()) + }) + + It("uses the active carrier when it cannot tell where a worker is, and leaves the verdict to the carrier", func() { + cmds := carrier.NewCommands(&cur) + cmds.UseWindow(window) + Expect(cmds.PingNode("unknown")).To(Succeed()) + Expect(next.commands.seen()).To(Equal([]string{"PingNode"})) + Expect(old.commands.seen()).To(BeEmpty()) + }) + + It("sends an agent request to the carrier that has agent workers", func() { + agents := carrier.NewAgents(&cur) + agents.UseWindow(window) + + agentsOn[cluster.CarrierTunnel] = true + _, err := agents.ExecuteMCPTool(ctx, mcpRemote.MCPToolRequest{}) + Expect(err).ToNot(HaveOccurred()) + Expect(next.agents.seen()).To(HaveLen(1)) + + agentsOn[cluster.CarrierTunnel] = false + agentsOn[cluster.CarrierNATS] = true + _, err = agents.DiscoverMCPTools(ctx, mcpRemote.MCPDiscoveryRequest{}) + Expect(err).ToNot(HaveOccurred()) + Expect(old.agents.seen()).To(HaveLen(1), "no agent worker has followed yet, so the old carrier still has them") + + agentsOn[cluster.CarrierNATS] = false + _, err = agents.ExecuteMCPTool(ctx, mcpRemote.MCPToolRequest{}) + Expect(err).ToNot(HaveOccurred()) + Expect(next.agents.seen()).To(HaveLen(2), "with none on either, the active carrier gives the error") + }) + + It("unloads a model through the previous carrier only for the nodes the active one could not reach", func() { + cmds := carrier.NewCommands(&cur) + cmds.UseWindow(window) + Expect(cmds.UnloadRemoteModel("m")).To(Succeed()) + Expect(next.commands.seen()).To(Equal([]string{"UnloadRemoteModelContext"})) + Expect(old.commands.seen()).To(BeEmpty(), "the active carrier did everything, so the previous one is not asked") + }) + + It("asks both carriers to delete the files of a model", func() { + cmds := carrier.NewCommands(&cur) + cmds.UseWindow(window) + Expect(cmds.DeleteModelFiles("m")).To(Succeed()) + Expect(next.commands.seen()).To(Equal([]string{"DeleteModelFiles"})) + Expect(old.commands.seen()).To(Equal([]string{"DeleteModelFiles"})) + }) + + It("goes back to the active set alone when it closes", func() { + attached["a"] = carrier.Attachment{NATS: true} + cmds := carrier.NewCommands(&cur) + cmds.UseWindow(window) + window.Close() + Expect(cmds.PingNode("a")).To(Succeed()) + Expect(old.commands.seen()).To(BeEmpty()) + Expect(next.commands.seen()).To(HaveLen(1)) + }) + }) +}) diff --git a/core/services/messaging/workqueue.go b/core/services/messaging/workqueue.go index 44df2e423e44..aa12e25ebb30 100644 --- a/core/services/messaging/workqueue.go +++ b/core/services/messaging/workqueue.go @@ -1,6 +1,16 @@ package messaging -import "context" +import ( + "context" + "errors" +) + +// ErrCarrierReleased is the cause that a context is cancelled with when the +// carrier its work ran on is released at the end of a drain. Work that sees it +// was cut off by the change of carrier, and not by its own caller or by a +// failure of its own, so it must not be run again on the carrier that took over: +// it did start, and the reaper decides what happens to its job. +var ErrCarrierReleased = errors.New("the carrier of this work was released") // WorkKind names a unit of competing-consumer work. The values match the claim // kinds of the self-hosted carrier so the two map one to one. From 227e0cd0b7fab58a04d7a4f905c0b411b6a476a4 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:33:37 +0000 Subject: [PATCH 109/162] feat(nodes): read the carriers of the workers, and keep a probe from reaping a healthy model A change of carrier asks the same questions about every worker: is it alive, which carriers is it attached to, and which can it follow. SwitchWorkers answers them from the node registry. A worker that reports its carriers is believed. One that reports nothing predates carrier switching, and it is on the tunnel if it holds one and on NATS if it does not. The columns for what a worker reports, and the call that stores them, are there for the heartbeat of a worker to fill. A worker that holds a tunnel has no address and binds its backends to loopback. Dialled directly, such an address names a port on the frontend, and the refused connection looks like a dead backend, so a forced change to NATS could make the health monitor and the reconciler reap models that run. The guarded client factory gives a worker that registered no address a client whose dial fails in the transport, which the code that decides about a backend already reads as "nothing learned". A worker that registered an address is dialled as before. The configuration no longer requires a NATS URL, because a deployment with only PostgreSQL runs on the tunnel. The three waits of a change of carrier can be set by flag. pgbus gets a probe of the LISTEN connection, for the report of what a replica could build, and the hint that tells replicas to read the row has its subject. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/run.go | 23 +++ core/config/distributed_config.go | 54 ++++- core/config/distributed_config_test.go | 34 ++++ core/services/carrier/pgbus.go | 9 + core/services/carrier/pgbus_test.go | 19 ++ core/services/messaging/subjects.go | 7 + core/services/nodes/follow_report.go | 236 ++++++++++++++++++++++ core/services/nodes/follow_report_test.go | 208 +++++++++++++++++++ core/services/nodes/registry.go | 23 ++- core/services/pgbus/listener.go | 23 +++ 10 files changed, 619 insertions(+), 17 deletions(-) create mode 100644 core/services/nodes/follow_report.go create mode 100644 core/services/nodes/follow_report_test.go diff --git a/core/cli/run.go b/core/cli/run.go index 0120a56c7cd6..3eb2342f6153 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -169,6 +169,9 @@ type RunCMD struct { PeerTLS bool `env:"LOCALAI_PEER_TLS" default:"false" help:"dial the other frontends over wss; they must sit behind TLS at the address they publish. Without it the peer link, which carries the credential of this replica and relayed requests, is clear text" group:"distributed"` PeerTLSCA string `env:"LOCALAI_PEER_TLS_CA" type:"existingfile" help:"PEM file with the CA that signs the certificate of the other frontends (system roots if empty); use with --peer-tls" group:"distributed"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` + CarrierPrepareTimeout string `env:"LOCALAI_CARRIER_PREPARE_TIMEOUT" help:"How long a change of carrier waits for every frontend to be ready before it is aborted (default 1m). The cluster setting switch.prepare_timeout, when stored, wins." group:"distributed"` + CarrierTransitionWindow string `env:"LOCALAI_CARRIER_TRANSITION_WINDOW" help:"How long a change of carrier waits for every frontend to confirm that it uses the new carrier (default 2m). The cluster setting switch.transition_window, when stored, wins." group:"distributed"` + CarrierMaxDrain string `env:"LOCALAI_CARRIER_MAX_DRAIN" help:"How long the previous carrier stays attached after a change of carrier, so that work that started on it can finish (default 15m). Work still running after that is failed by the reaper. The cluster setting switch.max_drain, when stored, wins." group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` StorageRegion string `env:"LOCALAI_STORAGE_REGION" default:"us-east-1" help:"S3 region" group:"distributed"` @@ -433,6 +436,26 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } opts = append(opts, config.WithNodeHeartbeatCheckpoint(d)) } + if r.CarrierPrepareTimeout != "" || r.CarrierTransitionWindow != "" || r.CarrierMaxDrain != "" { + var prepare, window, drain time.Duration + var err error + for _, f := range []struct { + env, value string + into *time.Duration + }{ + {"LOCALAI_CARRIER_PREPARE_TIMEOUT", r.CarrierPrepareTimeout, &prepare}, + {"LOCALAI_CARRIER_TRANSITION_WINDOW", r.CarrierTransitionWindow, &window}, + {"LOCALAI_CARRIER_MAX_DRAIN", r.CarrierMaxDrain, &drain}, + } { + if f.value == "" { + continue + } + if *f.into, err = parseDistributedDuration(f.env, f.value); err != nil { + return err + } + } + opts = append(opts, config.WithCarrierTimings(prepare, window, drain)) + } if r.ModelConfigResyncInterval != "" { d, err := parseDistributedDuration("LOCALAI_MODEL_CONFIG_RESYNC_INTERVAL", r.ModelConfigResyncInterval) if err != nil { diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 01a4c1e9850e..4d93e266032a 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -11,7 +11,8 @@ import ( ) // DistributedConfig holds configuration for horizontal scaling mode. -// When Enabled is true, PostgreSQL and NATS are required. +// When Enabled is true, PostgreSQL is required. NATS is needed only while the +// cluster runs on the NATS carrier. type DistributedConfig struct { Enabled bool // --distributed / LOCALAI_DISTRIBUTED InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) @@ -25,10 +26,20 @@ type DistributedConfig struct { // PeerTLSCA is a PEM file with the certificate authority that signs the // certificate of the other frontends. Empty uses the system roots. // LOCALAI_PEER_TLS_CA. - PeerTLSCA string - NatsURL string // --nats-url / LOCALAI_NATS_URL - StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) - RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) + PeerTLSCA string + // NatsURL is the NATS server of a deployment that uses NATS. A deployment with + // no URL and no row in the database runs on the tunnel carrier. The URL is + // copied into the cluster settings when none is stored there, and the cluster + // setting wins afterwards, so that every replica uses the same one. + NatsURL string // --nats-url / LOCALAI_NATS_URL + // CarrierPrepareTimeout, CarrierTransitionWindow and CarrierMaxDrain are the + // waits of a change of carrier on the replica that leads it, used when the + // cluster settings do not set them. Zero means the default. + CarrierPrepareTimeout time.Duration // --carrier-prepare-timeout + CarrierTransitionWindow time.Duration // --carrier-transition-window + CarrierMaxDrain time.Duration // --carrier-max-drain + StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) + RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) // RegistrationRequireAuth fails startup when distributed mode is enabled but // RegistrationToken is empty. The default (false) keeps the historical // fail-open behavior with a loud warning; production should set it so the @@ -153,9 +164,9 @@ func (c DistributedConfig) Validate() error { if !c.Enabled { return nil } - if c.NatsURL == "" { - return fmt.Errorf("distributed mode requires --nats-url / LOCALAI_NATS_URL") - } + // No NATS URL is required. The carrier is a setting of the cluster in the + // database: a new deployment with only PostgreSQL runs on the tunnel, and a + // deployment with a NATS URL stays on NATS. // S3 credentials must be paired if (c.StorageAccessKey != "" && c.StorageSecretKey == "") || (c.StorageAccessKey == "" && c.StorageSecretKey != "") { @@ -199,6 +210,15 @@ func (c DistributedConfig) Validate() error { if c.ModelConfigResyncInterval < 0 { return fmt.Errorf("%s must not be negative", FlagModelConfigResyncInterval) } + for name, d := range map[string]time.Duration{ + FlagCarrierPrepareTimeout: c.CarrierPrepareTimeout, + FlagCarrierTransitionWindow: c.CarrierTransitionWindow, + FlagCarrierMaxDrain: c.CarrierMaxDrain, + } { + if d < 0 { + return fmt.Errorf("%s must not be negative", name) + } + } return nil } @@ -474,6 +494,24 @@ const ( FlagDiskHeadroomCheck = "distributed-disk-headroom-check" ) +// Names of the timings of a change of carrier. +const ( + FlagCarrierPrepareTimeout = "carrier-prepare-timeout" + FlagCarrierTransitionWindow = "carrier-transition-window" + FlagCarrierMaxDrain = "carrier-max-drain" +) + +// WithCarrierTimings sets the waits of a change of carrier: how long prepare +// waits for every replica, how long commit waits for every replica to confirm, +// and how long the previous carrier stays attached. A zero keeps the default. +func WithCarrierTimings(prepareTimeout, transitionWindow, maxDrain time.Duration) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.CarrierPrepareTimeout = prepareTimeout + o.Distributed.CarrierTransitionWindow = transitionWindow + o.Distributed.CarrierMaxDrain = maxDrain + } +} + // FlagModelConfigResyncInterval names the model config resync interval. const FlagModelConfigResyncInterval = "model-config-resync-interval" diff --git a/core/config/distributed_config_test.go b/core/config/distributed_config_test.go index 4296054379a3..7192f1118f2e 100644 --- a/core/config/distributed_config_test.go +++ b/core/config/distributed_config_test.go @@ -89,9 +89,43 @@ var _ = Describe("DistributedConfig flag-name constants", func() { Entry("backend upgrade timeout", config.FlagBackendUpgradeTimeout, "backend-upgrade-timeout"), Entry("model load timeout", config.FlagModelLoadTimeout, "model-load-timeout"), Entry("model config resync interval", config.FlagModelConfigResyncInterval, "model-config-resync-interval"), + Entry("carrier prepare timeout", config.FlagCarrierPrepareTimeout, "carrier-prepare-timeout"), + Entry("carrier transition window", config.FlagCarrierTransitionWindow, "carrier-transition-window"), + Entry("carrier max drain", config.FlagCarrierMaxDrain, "carrier-max-drain"), ) }) +var _ = Describe("DistributedConfig without a NATS URL", func() { + It("is valid: a deployment with only PostgreSQL runs on the tunnel carrier", func() { + Expect(config.DistributedConfig{Enabled: true}.Validate()).To(Succeed()) + }) + + It("keeps the NATS URL when one is given", func() { + o := config.NewApplicationConfig(config.WithNatsURL("nats://localhost:4222")) + Expect(o.Distributed.NatsURL).To(Equal("nats://localhost:4222")) + }) + + DescribeTable("rejects a negative timing of a change of carrier", + func(set func(*config.DistributedConfig), flag string) { + c := config.DistributedConfig{Enabled: true} + set(&c) + err := c.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(flag)) + }, + Entry("prepare timeout", func(c *config.DistributedConfig) { c.CarrierPrepareTimeout = -time.Second }, config.FlagCarrierPrepareTimeout), + Entry("transition window", func(c *config.DistributedConfig) { c.CarrierTransitionWindow = -time.Second }, config.FlagCarrierTransitionWindow), + Entry("max drain", func(c *config.DistributedConfig) { c.CarrierMaxDrain = -time.Second }, config.FlagCarrierMaxDrain), + ) + + It("sets the timings with the options", func() { + o := config.NewApplicationConfig(config.WithCarrierTimings(10*time.Second, 20*time.Second, 30*time.Second)) + Expect(o.Distributed.CarrierPrepareTimeout).To(Equal(10 * time.Second)) + Expect(o.Distributed.CarrierTransitionWindow).To(Equal(20 * time.Second)) + Expect(o.Distributed.CarrierMaxDrain).To(Equal(30 * time.Second)) + }) +}) + var _ = Describe("DistributedConfig.Validate negative-duration errors", func() { It("rejects a negative BackendInstallTimeout with the flag name in the error", func() { c := config.DistributedConfig{ diff --git a/core/services/carrier/pgbus.go b/core/services/carrier/pgbus.go index 5da79e0327be..4b48b5b1a74f 100644 --- a/core/services/carrier/pgbus.go +++ b/core/services/carrier/pgbus.go @@ -44,6 +44,15 @@ type Fanout struct { Close func() } +// ProbeListen reports why a LISTEN connection to the database cannot be opened +// from this process, or nil when it can. It opens and closes one session. +func ProbeListen(ctx context.Context, dsn string) error { + if dsn == "" { + return errors.New("the connection string of the database is not known, so no LISTEN connection can be opened") + } + return pgbus.ProbeListen(ctx, dsn) +} + // NewPgbusFanout opens the LISTEN connection and returns the fan-out member. // // It is the one place that opens the connection. A deployment whose active diff --git a/core/services/carrier/pgbus_test.go b/core/services/carrier/pgbus_test.go index 186778bb734b..9b2bee888fd1 100644 --- a/core/services/carrier/pgbus_test.go +++ b/core/services/carrier/pgbus_test.go @@ -178,3 +178,22 @@ var _ = Describe("The pgbus fan-out of the tunnel carrier", func() { }) }) }) + +var _ = Describe("The probe of the LISTEN connection", func() { + It("succeeds against a database that accepts LISTEN, and leaves no session behind", func() { + db, dsn := testutil.SetupTestDBWithDSN() + Expect(carrier.ProbeListen(context.Background(), dsn)).To(Succeed()) + Eventually(func() int64 { return listeners(db) }, "5s").Should(BeZero()) + }) + + It("says why it cannot connect", func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := carrier.ProbeListen(ctx, "postgres://nobody@127.0.0.1:1/none?sslmode=disable&connect_timeout=1") + Expect(err).To(HaveOccurred()) + }) + + It("refuses an empty connection string with a reason that names the cause", func() { + Expect(carrier.ProbeListen(context.Background(), "")).To(MatchError(ContainSubstring("connection string"))) + }) +}) diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index e4139739eb94..55c160112d5e 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -20,6 +20,13 @@ const ( QueueWorkers = "workers" ) +// SubjectCarrierChanged is the hint that the cluster carrier row moved. The +// leader and the admin publish it after each move, on the carrier in use, and +// every replica listens on it. It carries the new row and is only a courtesy: a +// replica that never hears it reads the row at its next poll, and what the row +// says is what counts. +const SubjectCarrierChanged = "state.carrier" + // SubjectClaimWake is the hint that a unit of work was put in the claim queue. // A consumer that hears it looks for work at once and does not wait for its next // poll. The hint is a broadcast and broadcasts are at-most-once, so it can be diff --git a/core/services/nodes/follow_report.go b/core/services/nodes/follow_report.go new file mode 100644 index 000000000000..9493fab64886 --- /dev/null +++ b/core/services/nodes/follow_report.go @@ -0,0 +1,236 @@ +package nodes + +import ( + "context" + "errors" + "fmt" + "net" + "strings" + "sync" + "time" + + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/cluster" + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +// CarrierReport is what a worker says about itself in a heartbeat: the carriers +// it is attached to, the epoch of the cluster row it saw, the carriers it can +// attach to, and why it cannot attach to another one. +type CarrierReport struct { + Attached []cluster.Carrier + AttachedEpoch int64 + Follow []cluster.Carrier + FollowError string +} + +func joinCarriers(cs []cluster.Carrier) string { + out := make([]string, len(cs)) + for i, c := range cs { + out[i] = string(c) + } + return strings.Join(out, ",") +} + +func splitCarriers(s string) []cluster.Carrier { + if s == "" { + return nil + } + var out []cluster.Carrier + for _, part := range strings.Split(s, ",") { + if c := cluster.Carrier(strings.TrimSpace(part)); c == cluster.CarrierNATS || c == cluster.CarrierTunnel { + out = append(out, c) + } + } + return out +} + +// SetCarrierReport stores what a worker reported about its carriers. +func (r *NodeRegistry) SetCarrierReport(ctx context.Context, nodeID string, rep CarrierReport) error { + res := r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", nodeID).Updates(map[string]any{ + "attached": joinCarriers(rep.Attached), + "attached_epoch": rep.AttachedEpoch, + "follow": joinCarriers(rep.Follow), + "follow_error": rep.FollowError, + }) + if res.Error != nil { + return fmt.Errorf("storing the carrier report of node %s: %w", nodeID, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("storing the carrier report of node %s: %w", nodeID, gorm.ErrRecordNotFound) + } + return nil +} + +// PresenceReader says whether a worker holds a tunnel. *cluster.Registry is one. +type PresenceReader interface { + Presence(ctx context.Context, nodeID string, grace time.Duration) (cluster.Presence, error) +} + +// SwitchWorkers answers the questions that a change of carrier asks about the +// workers: which are alive, which carriers each is attached to, and which it can +// follow. A worker that reports its carriers is believed. One that does not is a +// worker that predates carrier switching: it holds a tunnel, or it is on NATS. +type SwitchWorkers struct { + reg *NodeRegistry + presence PresenceReader + grace time.Duration + stale time.Duration +} + +// NewSwitchWorkers returns the reader. A worker whose heartbeat is older than +// stale is not alive for this purpose; grace is the reconnect grace of the +// tunnel. +func NewSwitchWorkers(reg *NodeRegistry, presence PresenceReader, grace, stale time.Duration) *SwitchWorkers { + return &SwitchWorkers{reg: reg, presence: presence, grace: grace, stale: stale} +} + +var _ cluster.WorkerSource = (*SwitchWorkers)(nil) + +func (s *SwitchWorkers) attached(ctx context.Context, n *BackendNode) ([]cluster.Carrier, error) { + if reported := splitCarriers(n.Attached); len(reported) > 0 { + return reported, nil + } + p, err := s.presence.Presence(ctx, n.ID, s.grace) + if err != nil { + return nil, err + } + if p == cluster.PresenceConnected || p == cluster.PresenceReconnecting { + return []cluster.Carrier{cluster.CarrierTunnel}, nil + } + return []cluster.Carrier{cluster.CarrierNATS}, nil +} + +// Workers lists the workers that are alive. +func (s *SwitchWorkers) Workers(ctx context.Context) ([]cluster.WorkerInfo, error) { + list, err := s.reg.List(ctx) + if err != nil { + return nil, err + } + var out []cluster.WorkerInfo + for i := range list { + n := &list[i] + if n.Status == StatusPending || n.Status == StatusOffline || time.Since(n.LastHeartbeat) > s.stale { + continue + } + attached, err := s.attached(ctx, n) + if err != nil { + return nil, fmt.Errorf("finding the carriers of node %s: %w", n.Name, err) + } + follow := splitCarriers(n.Follow) + out = append(out, cluster.WorkerInfo{ + ID: n.ID, Name: n.Name, Type: n.NodeType, + Attached: attached, Reports: len(follow) > 0, Follow: follow, FollowError: n.FollowError, + }) + } + return out, nil +} + +// AttachedCarriers returns the carriers one worker is attached to. +func (s *SwitchWorkers) AttachedCarriers(ctx context.Context, nodeID string) ([]cluster.Carrier, error) { + n, err := s.reg.Get(ctx, nodeID) + if err != nil { + return nil, err + } + return s.attached(ctx, n) +} + +// AgentsAttached reports whether an agent worker that is alive is attached to c. +func (s *SwitchWorkers) AgentsAttached(ctx context.Context, c cluster.Carrier) (bool, error) { + list, err := s.reg.List(ctx) + if err != nil { + return false, err + } + for i := range list { + n := &list[i] + if n.NodeType != NodeTypeAgent || n.Status == StatusPending || n.Status == StatusOffline || time.Since(n.LastHeartbeat) > s.stale { + continue + } + attached, err := s.attached(ctx, n) + if err != nil { + return false, err + } + for _, a := range attached { + if a == c { + return true, nil + } + } + } + return false, nil +} + +// errTunnelWorkerLoopback is the failure of the dial that the guarded factory +// puts in the place of a dial that cannot mean anything. It is not the answer of +// a host about a backend, so nothing that decides whether a backend is dead may +// act on it. +var errTunnelWorkerLoopback = errors.New("the backend of this worker listens on its own loopback address and the worker holds no address that can be dialled, so it cannot be reached on this carrier") + +// guardTTL is how long the guarded factory keeps its answer about a worker. +const guardTTL = 10 * time.Second + +type guardedDirectFactory struct { + token string + direct BackendClientFactory + hasAddress func(nodeID string) (bool, error) + + mu sync.Mutex + known map[string]guardEntry +} + +type guardEntry struct { + at time.Time + has bool +} + +// NewGuardedDirectClientFactory is the direct client factory with one rule. A +// worker that holds a tunnel has no address and binds its backends to loopback. +// Dialled directly, a loopback address names a port on the frontend and not on the +// worker, and the refused connection that follows looks like a dead backend. The +// code that decides about a backend would then reap a model that runs. For such a +// worker the client gets a dial that fails in the transport, which every such +// code reads as "nothing learned". A worker that registered an address is dialled +// as before, so a deployment on NATS behaves as it did. +// +// hasAddress says whether the worker registered an address. When it cannot say, +// the client dials. +func NewGuardedDirectClientFactory(token string, hasAddress func(nodeID string) (bool, error)) BackendClientFactory { + return &guardedDirectFactory{token: token, direct: NewTokenClientFactory(token), hasAddress: hasAddress, known: map[string]guardEntry{}} +} + +func (f *guardedDirectFactory) NewClient(nodeID, address string, parallel bool) grpc.Backend { + if isLoopbackAddress(address) && f.withoutAddress(nodeID) { + return grpc.NewClientWithDialer(address, parallel, nil, false, f.token, + func(context.Context, string) (net.Conn, error) { return nil, errTunnelWorkerLoopback }) + } + return f.direct.NewClient(nodeID, address, parallel) +} + +func (f *guardedDirectFactory) withoutAddress(nodeID string) bool { + f.mu.Lock() + hit, ok := f.known[nodeID] + f.mu.Unlock() + if ok && time.Since(hit.at) < guardTTL { + return !hit.has + } + has, err := f.hasAddress(nodeID) + if err != nil { + return false + } + f.mu.Lock() + f.known[nodeID] = guardEntry{at: time.Now(), has: has} + f.mu.Unlock() + return !has +} + +func isLoopbackAddress(address string) bool { + host, _, err := net.SplitHostPort(address) + if err != nil { + host = address + } + if strings.EqualFold(host, "localhost") { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} diff --git a/core/services/nodes/follow_report_test.go b/core/services/nodes/follow_report_test.go new file mode 100644 index 000000000000..e74915666e2e --- /dev/null +++ b/core/services/nodes/follow_report_test.go @@ -0,0 +1,208 @@ +package nodes + +import ( + "context" + "errors" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +var _ = Describe("The workers as the switch of the carrier sees them", func() { + var ( + ctx context.Context + db *gorm.DB + reg *NodeRegistry + presence *fakePresence + workers *SwitchWorkers + ) + + register := func(name, nodeType, address string) *BackendNode { + GinkgoHelper() + n := &BackendNode{Name: name, NodeType: nodeType, Address: address, TokenHash: "h"} + Expect(reg.Register(ctx, n, true)).To(Succeed()) + return n + } + byName := func(list []cluster.WorkerInfo) map[string]cluster.WorkerInfo { + out := map[string]cluster.WorkerInfo{} + for _, w := range list { + out[w.Name] = w + } + return out + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + var err error + reg, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + presence = &fakePresence{by: map[string]cluster.Presence{}} + workers = NewSwitchWorkers(reg, presence, cluster.DefaultReconnectGrace, 5*time.Minute) + }) + + It("lists the workers that are alive and leaves out the ones that are not", func() { + register("live", NodeTypeBackend, "10.0.0.1:50051") + stale := register("stale", NodeTypeBackend, "10.0.0.2:50051") + offline := register("offline", NodeTypeBackend, "10.0.0.3:50051") + Expect(db.Model(&BackendNode{}).Where("id = ?", stale.ID).Update("last_heartbeat", time.Now().Add(-time.Hour)).Error).To(Succeed()) + Expect(db.Model(&BackendNode{}).Where("id = ?", offline.ID).Update("status", StatusOffline).Error).To(Succeed()) + pending := &BackendNode{Name: "pending", NodeType: NodeTypeBackend, TokenHash: "h"} + Expect(reg.Register(ctx, pending, false)).To(Succeed()) + + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(byName(list)).To(HaveLen(1)) + Expect(byName(list)).To(HaveKey("live")) + }) + + It("counts a worker that reports no capabilities as one that cannot follow, and says it is on NATS when it holds no tunnel", func() { + register("old", NodeTypeBackend, "10.0.0.1:50051") + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + w := byName(list)["old"] + Expect(w.Reports).To(BeFalse()) + Expect(w.Attached).To(Equal([]cluster.Carrier{cluster.CarrierNATS})) + }) + + It("says a worker that holds a tunnel is attached to the tunnel", func() { + n := register("tun", NodeTypeBackend, "") + presence.by[n.ID] = cluster.PresenceConnected + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(byName(list)["tun"].Attached).To(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + }) + + It("says a worker that was attached to the tunnel a moment ago is still attached to it", func() { + n := register("blink", NodeTypeBackend, "") + presence.by[n.ID] = cluster.PresenceReconnecting + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(byName(list)["blink"].Attached).To(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + }) + + It("takes what a worker reports over what can be inferred", func() { + n := register("both", NodeTypeBackend, "10.0.0.1:50051") + presence.by[n.ID] = cluster.PresenceConnected + Expect(reg.SetCarrierReport(ctx, n.ID, CarrierReport{ + Attached: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, AttachedEpoch: 7, + Follow: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, FollowError: "", + })).To(Succeed()) + + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + w := byName(list)["both"] + Expect(w.Reports).To(BeTrue()) + Expect(w.Attached).To(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + Expect(w.Follow).To(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + }) + + It("carries the reason a worker gives for not following", func() { + n := register("tunnel-only", NodeTypeBackend, "") + Expect(reg.SetCarrierReport(ctx, n.ID, CarrierReport{ + Attached: []cluster.Carrier{cluster.CarrierTunnel}, Follow: []cluster.Carrier{cluster.CarrierTunnel}, + FollowError: "it has no address that the frontends can reach", + })).To(Succeed()) + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(byName(list)["tunnel-only"].FollowError).To(Equal("it has no address that the frontends can reach")) + }) + + It("answers where one worker is attached", func() { + n := register("w", NodeTypeBackend, "10.0.0.1:50051") + a, err := workers.AttachedCarriers(ctx, n.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(a).To(Equal([]cluster.Carrier{cluster.CarrierNATS})) + + _, err = workers.AttachedCarriers(ctx, "nobody") + Expect(err).To(HaveOccurred()) + }) + + It("says whether an agent worker is attached to a carrier", func() { + register("backend", NodeTypeBackend, "10.0.0.1:50051") + agent := register("agent", NodeTypeAgent, "") + presence.by[agent.ID] = cluster.PresenceConnected + + onTunnel, err := workers.AgentsAttached(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(BeTrue()) + onNATS, err := workers.AgentsAttached(ctx, cluster.CarrierNATS) + Expect(err).ToNot(HaveOccurred()) + Expect(onNATS).To(BeFalse(), "the backend worker is on NATS, and it is not an agent worker") + }) + + It("returns the error of the presence read and not a guess", func() { + n := register("w", NodeTypeBackend, "") + presence.err = errors.New("the database is down") + _, err := workers.AttachedCarriers(ctx, n.ID) + Expect(err).To(HaveOccurred()) + }) +}) + +var _ = Describe("The direct client factory for a worker that holds a tunnel", func() { + var ( + ctx context.Context + known map[string]bool + asked int + factory BackendClientFactory + ) + + BeforeEach(func() { + ctx = context.Background() + known = map[string]bool{"tunnel-only": false, "routable": true} + asked = 0 + factory = NewGuardedDirectClientFactory("tok", func(nodeID string) (bool, error) { + asked++ + has, ok := known[nodeID] + if !ok { + return false, errors.New("unknown node") + } + return has, nil + }) + }) + + health := func(nodeID, address string) (grpc.Backend, bool, error) { + c := factory.NewClient(nodeID, address, false) + hc, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + ok, err := c.HealthCheck(hc) + return c, ok, err + } + + It("does not dial the loopback address of a worker that registered no address, and reports a failure of the transport", func() { + c, ok, err := health("tunnel-only", "127.0.0.1:1") + Expect(ok).To(BeFalse()) + Expect(err).To(HaveOccurred()) + Expect(grpc.TransportFailureOf(c)).To(HaveOccurred(), "a caller that decides whether a backend is dead must not count this") + Expect(classifyProbeOutcome(ok, err)).ToNot(Equal(ProbeAlive)) + }) + + It("dials the loopback address of a worker that registered an address, as before", func() { + c, ok, _ := health("routable", "127.0.0.1:1") + Expect(ok).To(BeFalse()) + Expect(grpc.TransportFailureOf(c)).ToNot(HaveOccurred(), "a refused connection on a direct dial is the answer of the host about the backend") + }) + + It("does not ask the registry for an address that is not on loopback", func() { + _, _, _ = health("tunnel-only", "10.255.255.1:1") + Expect(asked).To(BeZero()) + }) + + It("dials when the registry cannot say", func() { + c, _, _ := health("unknown", "127.0.0.1:1") + Expect(grpc.TransportFailureOf(c)).ToNot(HaveOccurred()) + }) + + It("asks the registry once for a burst of clients of one worker", func() { + for range 5 { + factory.NewClient("tunnel-only", "127.0.0.1:1", false) + } + Expect(asked).To(Equal(1)) + }) +}) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index a20bf11b2241..bd5b0e029051 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -27,11 +27,11 @@ type BackendNode struct { WorkerIncarnation string `gorm:"size:36" json:"worker_incarnation,omitempty"` ID string `gorm:"primaryKey;size:36" json:"id"` Name string `gorm:"uniqueIndex;size:255" json:"name"` - NodeType string `gorm:"size:32;default:backend" json:"node_type"` // backend, agent - Address string `gorm:"size:255" json:"address"` // host:port for gRPC - HTTPAddress string `gorm:"size:255" json:"http_address"` // host:port for HTTP file transfer - Status string `gorm:"size:32;default:registering" json:"status"` // registering, healthy, unhealthy, draining, pending - TokenHash string `gorm:"size:64" json:"-"` // SHA-256 of registration token + NodeType string `gorm:"size:32;default:backend" json:"node_type"` // backend, agent + Address string `gorm:"size:255" json:"address"` // host:port for gRPC + HTTPAddress string `gorm:"size:255" json:"http_address"` // host:port for HTTP file transfer + Status string `gorm:"size:32;default:registering" json:"status"` // registering, healthy, unhealthy, draining, pending + TokenHash string `gorm:"size:64" json:"-"` // SHA-256 of registration token // TunnelTokenHash is the SHA-256 of the own credential of this node for the // tunnel: the token that it presents at GET /api/cluster/connect. It is not // the registration token. Registration mints a new random secret for each @@ -43,8 +43,8 @@ type BackendNode struct { // tunnel until it registers again. The column cannot be filled in, because // the secret exists only in the response that minted it. TunnelTokenHash string `gorm:"size:64" json:"-"` - TotalVRAM uint64 `gorm:"column:total_vram" json:"total_vram"` // Total GPU VRAM in bytes - AvailableVRAM uint64 `gorm:"column:available_vram" json:"available_vram"` // Available GPU VRAM in bytes + TotalVRAM uint64 `gorm:"column:total_vram" json:"total_vram"` // Total GPU VRAM in bytes + AvailableVRAM uint64 `gorm:"column:available_vram" json:"available_vram"` // Available GPU VRAM in bytes // ReservedVRAM is a soft, in-tick reservation deducted by the scheduler when // it picks this node to load a model. Workers reset it back to 0 on each // heartbeat (the worker is the source of truth for actual free VRAM); the @@ -113,8 +113,13 @@ type BackendNode struct { // carrier row that the worker saw when it reported. During a change of // carrier a worker is attached to both, and the frontends route by this. // Both are empty for a worker that does not report them. - Attached string `gorm:"column:attached;size:64" json:"attached,omitempty"` - AttachedEpoch int64 `gorm:"column:attached_epoch;default:0" json:"attached_epoch,omitempty"` + Attached string `gorm:"column:attached;size:64" json:"attached,omitempty"` + AttachedEpoch int64 `gorm:"column:attached_epoch;default:0" json:"attached_epoch,omitempty"` + // Follow lists the carriers the worker reported it can attach to, separated by + // commas, and FollowError is why it cannot attach to another one, in its own + // words. A worker that predates carrier switching reports neither. + Follow string `gorm:"column:follow;size:64" json:"follow,omitempty"` + FollowError string `gorm:"column:follow_error;size:255" json:"follow_error,omitempty"` APIKeyID string `gorm:"size:36" json:"-"` // auto-provisioned API key ID (for cleanup) AuthUserID string `gorm:"size:36" json:"-"` // auto-provisioned user ID (for cleanup) LastHeartbeat time.Time `gorm:"column:last_heartbeat" json:"last_heartbeat"` diff --git a/core/services/pgbus/listener.go b/core/services/pgbus/listener.go index f656401d85bf..3ce42d005c48 100644 --- a/core/services/pgbus/listener.go +++ b/core/services/pgbus/listener.go @@ -153,6 +153,29 @@ func dial(ctx context.Context, dsn, appName string) (*pgx.Conn, error) { return pgx.ConnectConfig(ctx, connCfg) } +// ProbeListen opens a LISTEN session, listens on a channel, and closes the +// session. It says whether this process could run a Bus against the database +// now: a pooler in transaction mode, a firewall, or a role that may not LISTEN +// each fail here and not at the first subscription. It writes nothing. +func ProbeListen(ctx context.Context, dsn string) error { + if dsn == "" { + return errors.New("pgbus: no connection string for the LISTEN connection") + } + conn, err := dial(ctx, dsn, applicationNamePrefix+"probe") + if err != nil { + return err + } + defer func() { + closeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Second) + defer cancel() + _ = conn.Close(closeCtx) + }() + if _, err := conn.Exec(ctx, "LISTEN "+pgx.Identifier{"localai_probe"}.Sanitize()); err != nil { + return fmt.Errorf("pgbus: the database refused LISTEN: %w", err) + } + return nil +} + // command passes a LISTEN or UNLISTEN to the goroutine that owns the connection // and waits for it. To return before the server acknowledged it would lose every // message published in the gap. From a333b4a5911760abb29f459e199646c1a3f80063 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:36:12 +0000 Subject: [PATCH 110/162] feat(jobs): complete the claim of a run that the end of a drain cut off When a change of carrier ends its drain, the claim loop of the old carrier is stopped, and a run that is still going is cancelled with messaging.ErrCarrierReleased. The driver used to read that as a link that broke, release the claim, and let the carrier that took over start the run again. The run had started, and the worker's run dies with its request, so a second start is a second execution. The driver now completes the claim and leaves the job as it is. The reaper fails a job that stays running, as it does for a job that lost its worker. Any other reason for the end of the context, such as the stop of the replica, still releases the claim. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/jobs/claim_driver.go | 21 +++++++++++ core/services/jobs/dispatch_loop_test.go | 47 ++++++++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/core/services/jobs/claim_driver.go b/core/services/jobs/claim_driver.go index d502641b53d8..3bbee9ee9dfb 100644 --- a/core/services/jobs/claim_driver.go +++ b/core/services/jobs/claim_driver.go @@ -191,6 +191,9 @@ func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb s for range maxDrivePicks { nodeID, nodeType, err := d.cfg.Picker.PickConnectedExcluding(ctx, tried) if err != nil { + if carrierReleased(ctx) { + return d.cutOff(jobID) + } if last != nil { // Every worker that was offered the run had no slot. That is the // evidence, and the end of the list is not. @@ -206,6 +209,8 @@ func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb s switch { case err == nil: return d.persistTerminal(&reply) + case carrierReleased(ctx): + return d.cutOff(jobID) case errors.Is(err, workerctl.ErrWorkerBusy): last = fmt.Errorf("offering a %s claim to agent worker %q: %w", kind, nodeID, err) continue @@ -221,6 +226,22 @@ func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb s return last } +// carrierReleased reports that ctx ended because the carrier of the loop was +// released at the end of a drain. It is not a stop of the replica, and not the +// cancel of a job. +func carrierReleased(ctx context.Context) bool { + return errors.Is(context.Cause(ctx), messaging.ErrCarrierReleased) +} + +// cutOff settles the claim of a run that the end of the drain cut off. The claim +// is completed, and not released: the run did start, so a carrier that took over +// the queue must not start it a second time. The job is left as it is, and the +// reaper fails it if it stays running, as it does for a job that lost its worker. +func (d *AgentDriver) cutOff(jobID string) error { + xlog.Warn("A run was cut off by the end of the drain of its carrier; its job is left to the reaper", "job", jobID) + return nil +} + // rebroadcast publishes the broadcast that a progress line asked for. A line with // no subject is meant for the caller alone, as every private tick is, and is not // passed on: the allow list denies the empty subject, and passing it would log a diff --git a/core/services/jobs/dispatch_loop_test.go b/core/services/jobs/dispatch_loop_test.go index 0f6e19f38b28..630b6ff6ad35 100644 --- a/core/services/jobs/dispatch_loop_test.go +++ b/core/services/jobs/dispatch_loop_test.go @@ -120,6 +120,53 @@ var _ = Describe("The dispatch loop", func() { Expect(row.ClaimedBy).To(Equal("replica-a")) }) + Describe("when the carrier of the loop is released", func() { + // run starts a loop whose context the spec can end with a cause, and waits + // until a run is in flight on it. + run := func() (cancel context.CancelCauseFunc, loop *DispatchLoop) { + GinkgoHelper() + started := make(chan struct{}, 1) + control.do = func(ctx context.Context, _, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + started <- struct{}{} + <-ctx.Done() + return workerctl.RunReply{}, ctx.Err() + } + live(db, "replica-a") + loopCtx, cancel := context.WithCancelCause(ctx) + var err error + loop, err = NewDispatchLoop(DispatchConfig{ + DB: db, Owner: "replica-a", Picker: picker, Control: control, Broadcast: broadcast, Store: store, Hints: bus, + Interval: 50 * time.Millisecond, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(loop.Start(loopCtx)).To(Succeed()) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, ci("j1"))).To(Succeed()) + Eventually(started, 10*time.Second).Should(Receive()) + return cancel, loop + } + + It("completes the claim of a run that the end of the drain cut off, and leaves its job to the reaper, so the run is not offered again", func() { + cancel, loop := run() + cancel(messaging.ErrCarrierReleased) + loop.Stop() + + Expect(rows()).To(BeZero(), "the run did start, and the carrier that took over must not start it again") + Expect(store.terminals()).To(BeEmpty(), "the reaper decides about the job") + Expect(control.seen()).To(HaveLen(1)) + }) + + It("still releases the claim of a run that ends for another reason, such as the stop of the replica", func() { + cancel, loop := run() + cancel(context.Canceled) + loop.Stop() + + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed()) + Expect(row.State).To(Equal(ClaimPending)) + Expect(row.Attempts).To(Equal(1)) + }) + }) + It("ends the run of a job that is cancelled, closes the job as cancelled and removes the claim", func() { started := make(chan struct{}) control.do = func(ctx context.Context, _, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { From 519cd95fdf9daf63f9fe5a3e4bbdff82151102dd Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:42:02 +0000 Subject: [PATCH 111/162] feat(application): start on the carrier the cluster names, and follow its changes A frontend no longer refuses a cluster that runs on the tunnel. It reads the row, seeds it when no replica has (the tunnel for a deployment with only PostgreSQL, NATS for one with a NATS URL, which keeps the behaviour of an existing deployment), and builds the set the row names. For the tunnel that is pgbus, the claim queue, the control client and the dialers behind the holders. For NATS it is what it was. A NATS URL that cannot be parsed stops the start, and one whose server is not up does not. The NATS URL becomes a setting of the cluster. The flag is copied there when none is stored, and a stored URL wins, so that every replica uses the same one. While the tunnel is the carrier in use or is draining, each replica claims queued work and drives it on agent workers, through the dialer of the tunnel set and not through the holder, so that a run that started on the tunnel finishes there. When the tunnel is released, the rows that are still pending are taken out of the table and published to the queue of the carrier that took over. Every replica now follows the row: a swapper that polls it and wakes on a hint, a leader loop that drives the protocol under an advisory lock, and a report of which carriers the replica could build. The report is refreshed when a dry run asks for it. The routing of calls to workers during a change is wired into every holder and into the dialer, and the draining carrier forgets a node when it is removed. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/carrier_runtime.go | 387 +++++++++++++++++++ core/application/distributed.go | 334 +++++++++++----- core/application/distributed_carrier_test.go | 106 ++++- core/application/startup.go | 4 + core/services/advisorylock/keys.go | 21 +- core/services/carrier/nats.go | 13 +- core/services/messaging/subjects.go | 6 + 7 files changed, 759 insertions(+), 112 deletions(-) create mode 100644 core/application/carrier_runtime.go diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go new file mode 100644 index 000000000000..02de689c0af0 --- /dev/null +++ b/core/application/carrier_runtime.go @@ -0,0 +1,387 @@ +package application + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sync/atomic" + "time" + + "github.com/mudler/xlog" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/jobs" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/pkg/sanitize" +) + +const ( + // natsConnectWait is how long a replica waits for the NATS server when it + // prepares the carrier for a change. At start-up it does not wait, because a + // frontend may start before its broker. + natsConnectWait = 10 * time.Second + // natsProbeWait is how long the report of what a replica can build waits for + // the NATS server. + natsProbeWait = 3 * time.Second + // availabilityEvery is how often a replica writes what it can build without + // being asked. A dry run asks, so this only keeps the status view from + // showing nothing. + availabilityEvery = 5 * time.Minute + // availabilityMaxAge is how old a report of what a replica can build may be + // for a change of carrier to rely on it. A dry run asks the replicas to look + // again first, so a report older than this is from a replica that did not + // answer. + availabilityMaxAge = 30 * time.Second + // probeWait is how long a dry run waits for the replicas to answer. + probeWait = 5 * time.Second + // hintPublishTimeout bounds the publish of a hint. + hintPublishTimeout = 3 * time.Second +) + +// carrierRuntime builds the carrier sets of this replica and reports what it +// could build. It is the one place of the application that names the carriers: +// everything above it holds the holders. +type carrierRuntime struct { + cfg *config.ApplicationConfig + db *gorm.DB + settings *cluster.SettingsStore + registry *nodes.NodeRegistry + fileMgr *storage.FileManager + instanceID string + jobStore *jobs.JobStore + clusterReg *cluster.Registry + tunnels *tunnel.Registry + peerPool *tunnel.PeerPool + + // bus is the broadcaster holder. The work that a set starts uses it, so that + // its subscriptions follow the sets. It is set before any set starts. + bus *carrier.Broadcaster + // active is the pointer the holders read. + active *atomic.Pointer[carrier.Set] + + probeWake chan struct{} +} + +// natsURL returns the NATS server of the cluster: the stored setting, else the +// flag of this replica. +func (rt *carrierRuntime) natsURL(ctx context.Context) (string, error) { + v, ok, err := rt.settings.Get(ctx, cluster.SettingNATSURL) + if err != nil { + return "", err + } + if ok && v != "" { + return v, nil + } + return rt.cfg.Distributed.NatsURL, nil +} + +// Build builds the set of a carrier for a change. It waits for the carrier to be +// usable, because a replica that reports ready must be. +func (rt *carrierRuntime) Build(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*carrier.Set, error) { + return rt.build(ctx, row, target, true) +} + +func (rt *carrierRuntime) build(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier, wait bool) (*carrier.Set, error) { + switch target { + case cluster.CarrierNATS: + return rt.buildNATS(ctx, row, wait) + case cluster.CarrierTunnel: + return rt.buildTunnel(ctx, row) + } + return nil, fmt.Errorf("%w: %q", cluster.ErrInvalidCarrier, target) +} + +func (rt *carrierRuntime) natsClient(ctx context.Context, wait time.Duration) (*messaging.Client, string, error) { + url, err := rt.natsURL(ctx) + if err != nil { + return nil, "", err + } + if url == "" { + return nil, "", errors.New("no NATS URL is configured: store one in the cluster settings (nats.url) or start the frontend with --nats-url") + } + d := rt.cfg.Distributed + auth := d.NatsAuthConfig() + if auth.RequireAuth && (auth.ServiceUserJWT == "" || auth.ServiceUserSeed == "") { + return nil, "", errors.New("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") + } + // A URL or a credential the client cannot use stops here. A server that is + // not up does not, unless the caller waits for it: the client keeps trying. + client, err := messaging.New(url, d.NatsMessagingOptions("", "")...) + if err != nil { + return nil, "", fmt.Errorf("connecting to NATS: %w", err) + } + if wait > 0 { + deadline := time.Now().Add(wait) + for !client.IsConnected() { + if time.Now().After(deadline) || ctx.Err() != nil { + client.Close() + return nil, "", fmt.Errorf("the NATS server at %s did not answer within %s: check the address and the credentials", sanitize.URL(url), wait) + } + select { + case <-ctx.Done(): + case <-time.After(50 * time.Millisecond): + } + } + } + return client, url, nil +} + +func (rt *carrierRuntime) buildNATS(ctx context.Context, row cluster.CarrierRow, wait bool) (*carrier.Set, error) { + var waitFor time.Duration + if wait { + waitFor = natsConnectWait + } + client, url, err := rt.natsClient(ctx, waitFor) + if err != nil { + return nil, err + } + d := rt.cfg.Distributed + set, err := carrier.NewNATSSet(carrier.NATSOptions{ + Client: client, + Epoch: row.Epoch, + Registry: rt.registry, + InstallTimeout: d.BackendInstallTimeoutOrDefault(), + UpgradeTimeout: d.BackendUpgradeTimeoutOrDefault(), + Token: d.RegistrationToken, + S3Staging: d.StorageURL != "", + FileManager: rt.fileMgr, + HTTPAddrFor: func(nodeID string) (string, error) { + node, err := rt.registry.Get(context.Background(), nodeID) + if err != nil { + return "", err + } + if node.HTTPAddress == "" { + return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) + } + return node.HTTPAddress, nil + }, + NodeHasAddress: func(nodeID string) (bool, error) { + node, err := rt.registry.Get(context.Background(), nodeID) + if err != nil { + return false, err + } + return node.Address != "", nil + }, + }) + if err != nil { + client.Close() + return nil, err + } + xlog.Info("Connected to NATS", "url", sanitize.URL(url), "epoch", row.Epoch) + return set, nil +} + +func (rt *carrierRuntime) buildTunnel(ctx context.Context, row cluster.CarrierRow) (*carrier.Set, error) { + d := rt.cfg.Distributed + fanout, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: rt.db, DSN: rt.cfg.Auth.DatabaseURL}) + if err != nil { + return nil, err + } + selector := nodes.NewAgentSelector(rt.registry, rt.clusterReg, rt.instanceID) + set, err := carrier.NewTunnelSet(carrier.TunnelOptions{ + Epoch: row.Epoch, + Dialer: tunnel.NewWorkerDialer(rt.tunnels, rt.peerPool), + Fanout: fanout, + WorkQueue: jobs.NewClaimQueue(rt.db, fanout.Broadcaster), + AgentSelector: selector, + Registry: rt.registry, + InstallTimeout: d.BackendInstallTimeoutOrDefault(), + UpgradeTimeout: d.BackendUpgradeTimeoutOrDefault(), + Token: d.RegistrationToken, + S3Staging: d.StorageURL != "", + FileManager: rt.fileMgr, + }) + if err != nil { + fanout.Close() + return nil, err + } + + // While the tunnel is the carrier in use, and while it drains, this replica + // claims queued work and drives it on agent workers. The loop dials through + // the dialer of this set and not through the holder: a run that started on the + // tunnel finishes there. + set.Start = func(ctx context.Context) (func(), error) { + loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{ + DB: rt.db, + Owner: rt.instanceID, + Picker: selector, + Control: nodes.NewControlClient(set.Dialer, d.RegistrationToken), + Broadcast: nodes.NewRebroadcaster(rt.bus), + Store: rt.jobStore, + Hints: rt.bus, + }) + if err != nil { + return nil, err + } + if err := loop.Start(ctx); err != nil { + return nil, err + } + xlog.Info("This replica claims queued work") + return loop.Stop, nil + } + + // When the tunnel is released, what is still queued moves to the queue of the + // carrier that took over. The rows are taken out of the table first and then + // published, so a job is lost, and not run twice, if this replica dies in + // between. A job that is lost stays running until the reaper fails it. + set.Handoff = func(ctx context.Context, next *carrier.Set) error { + rows, err := jobs.MigratePending(ctx, rt.db) + if err != nil { + return err + } + var failed error + moved := 0 + for _, r := range rows { + if err := next.WorkQueue.Enqueue(ctx, messaging.WorkKind(r.Kind), json.RawMessage(r.Payload)); err != nil { + failed = errors.Join(failed, fmt.Errorf("publishing the %s unit %s: %w", r.Kind, r.ID, err)) + continue + } + moved++ + } + if len(rows) > 0 { + xlog.Info("Queued work moved to the carrier that took over", "carrier", next.Name, "moved", moved, "lost", len(rows)-moved) + } + return failed + } + return set, nil +} + +// availability says, for each carrier, why this replica could not build it now. +func (rt *carrierRuntime) availability(ctx context.Context) map[cluster.Carrier]string { + out := map[cluster.Carrier]string{} + cur := rt.active.Load() + + switch { + case cur != nil && cur.Name == cluster.CarrierNATS: + out[cluster.CarrierNATS] = "" + default: + client, _, err := rt.natsClient(ctx, natsProbeWait) + if err != nil { + out[cluster.CarrierNATS] = err.Error() + } else { + client.Close() + out[cluster.CarrierNATS] = "" + } + } + + if cur != nil && cur.Name == cluster.CarrierTunnel { + out[cluster.CarrierTunnel] = "" + } else if err := carrier.ProbeListen(ctx, rt.cfg.Auth.DatabaseURL); err != nil { + out[cluster.CarrierTunnel] = err.Error() + } else { + out[cluster.CarrierTunnel] = "" + } + return out +} + +// reportAvailability looks at what this replica can build and writes it. +func (rt *carrierRuntime) reportAvailability(ctx context.Context) { + probeCtx, cancel := context.WithTimeout(ctx, 20*time.Second) + defer cancel() + if err := rt.clusterReg.ReportAvailability(ctx, rt.instanceID, rt.availability(probeCtx)); err != nil && ctx.Err() == nil { + xlog.Warn("Could not report which carriers this replica can build", "error", err) + } +} + +// RunAvailability reports what this replica can build at start, then when asked +// and every availabilityEvery. +func (rt *carrierRuntime) RunAvailability(ctx context.Context) { + rt.reportAvailability(ctx) + ticker := time.NewTicker(availabilityEvery) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + case <-rt.probeWake: + } + rt.reportAvailability(ctx) + } +} + +// wakeProbe asks RunAvailability to look now. +func (rt *carrierRuntime) wakeProbe() { + select { + case rt.probeWake <- struct{}{}: + default: + } +} + +// ProbeReplicas asks every replica to look at what it can build, and waits until +// the live replicas have answered or probeWait has passed. A replica that did not +// answer shows up as stale in the preflight. +func (rt *carrierRuntime) ProbeReplicas(ctx context.Context) { + asked := time.Now() + rt.wakeProbe() + pubCtx, cancel := context.WithTimeout(ctx, hintPublishTimeout) + defer cancel() + done := make(chan struct{}) + go func() { + defer close(done) + if err := rt.bus.Publish(messaging.SubjectCarrierProbe, map[string]string{"from": rt.instanceID}); err != nil { + xlog.Debug("Could not ask the other replicas to look at the carriers", "error", err) + } + }() + select { + case <-done: + case <-pubCtx.Done(): + } + + deadline := asked.Add(probeWait) + for time.Now().Before(deadline) && ctx.Err() == nil { + live, err := rt.clusterReg.ListLive(ctx, cluster.InstanceLiveness) + if err == nil { + fresh := true + for _, in := range live { + if _, known := in.AvailabilityFor(cluster.CarrierNATS); !known || in.AvailabilityAge > time.Since(asked) { + fresh = false + break + } + } + if fresh { + return + } + } + select { + case <-ctx.Done(): + case <-time.After(150 * time.Millisecond): + } + } +} + +// workCounts counts the work that a change of carrier can touch. +type workCounts struct{ db *gorm.DB } + +func (w workCounts) InFlight(ctx context.Context) (cluster.InFlight, error) { + var out cluster.InFlight + count := func(dst *int, q *gorm.DB) error { + var n int64 + if err := q.Count(&n).Error; err != nil { + return err + } + *dst = int(n) + return nil + } + db := w.db.WithContext(ctx) + if err := count(&out.Loads, db.Model(&nodes.NodeModel{}).Where("state IN ?", []string{"loading", "staging"})); err != nil { + return out, fmt.Errorf("counting the models that load: %w", err) + } + if err := count(&out.Jobs, db.Model(&jobs.JobRecord{}).Where("status = ?", "running")); err != nil { + return out, fmt.Errorf("counting the jobs that run: %w", err) + } + if err := count(&out.PendingClaims, db.Model(&jobs.WorkClaim{}).Where("state = ?", jobs.ClaimPending)); err != nil { + return out, fmt.Errorf("counting the pending claims: %w", err) + } + if err := count(&out.ClaimedClaims, db.Model(&jobs.WorkClaim{}).Where("state = ?", jobs.ClaimClaimed)); err != nil { + return out, fmt.Errorf("counting the claimed claims: %w", err) + } + return out, nil +} diff --git a/core/application/distributed.go b/core/application/distributed.go index 1f368ff92605..2660d165a6da 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -98,8 +98,24 @@ type DistributedServices struct { // dial seams on it; nothing calls it while the cluster runs on NATS. WorkerDialer *tunnel.WorkerDialer + // Settings holds the settings that decide what the whole cluster does: the + // NATS addresses and the waits of a change of carrier. + Settings *cluster.SettingsStore + // Switch is the protocol of a change of carrier. The admin API calls it, and + // the leader drives it. + Switch *cluster.Switch + // Swapper is what this replica does when the cluster changes its carrier. + Swapper *carrier.Swapper + // Runtime builds the carrier sets of this replica and reports what it could + // build. + Runtime *carrierRuntime + // active names the carrier set the holders forward to. active *atomic.Pointer[carrier.Set] + // window routes the calls to workers while two carriers are attached. + window *carrier.Window + // workers answers the questions about workers that a change asks. + workers *nodes.SwitchWorkers shutdownOnce sync.Once } @@ -114,6 +130,12 @@ func (ds *DistributedServices) Disconnect(nodeID string) bool { set.ForgetNode(nodeID) } } + // The carrier that is draining caches the same per-node state. + if ds.window != nil { + if prev := ds.window.Previous(); prev != nil && prev.ForgetNode != nil { + prev.ForgetNode(nodeID) + } + } return held } @@ -124,7 +146,12 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { - // The tunnels first: their claims name this replica, and a peer that + // Stop following the cluster carrier and stop the work that a carrier started, + // such as the claims this replica drives, before anything is torn down. + if ds.Swapper != nil { + ds.Swapper.Close() + } + // The tunnels next: their claims name this replica, and a peer that // asks who owns a worker must not be told a replica that is leaving. if ds.Tunnels != nil { ds.Tunnels.Close() @@ -195,47 +222,31 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("distributed mode requires auth database to be initialized first") } - // The cluster carrier row says which transport the whole cluster uses. - // The first replica to start writes it; later replicas, and every restart, - // follow it. Validate above has already refused a missing NATS URL, so a - // replica that seeds the row seeds NATS. + // The cluster carrier row says which transport the whole cluster uses. The + // first replica to start writes it; later replicas, and every restart, follow + // it and never their own flags. A deployment that has a NATS URL (the flag, or + // a URL stored in the cluster settings) is seeded on NATS, so an existing + // deployment keeps what it has. One that has only PostgreSQL is seeded on the + // tunnel. carrierStore, err := cluster.NewCarrierStore(authDB) if err != nil { return nil, fmt.Errorf("initializing cluster carrier state: %w", err) } - carrierRow, seeded, err := carrierStore.Seed(context.Background(), cluster.CarrierNATS, cfg.Distributed.InstanceID) + settingsStore, err := cluster.NewSettingsStore(authDB) if err != nil { - return nil, fmt.Errorf("reading cluster carrier state: %w", err) - } - if seeded { - xlog.Info("Cluster carrier seeded", "carrier", carrierRow.Active, "epoch", carrierRow.Epoch) - } - if carrierRow.Active != cluster.CarrierNATS { - // No silent fallback to the carrier this replica has flags for: a - // cluster that has moved to another carrier would split in two. - return nil, fmt.Errorf("the cluster carrier is %q (epoch %d, set by %s), and this build runs the %q carrier only", - carrierRow.Active, carrierRow.Epoch, carrierRow.ChangedBy, cluster.CarrierNATS) - } - - // Connect to NATS. A URL or credential the client cannot use stops the - // start. A server that is not up yet does not: the client keeps retrying, - // so a frontend can start before its broker. - natsAuth := cfg.Distributed.NatsAuthConfig() - if natsAuth.RequireAuth && (natsAuth.ServiceUserJWT == "" || natsAuth.ServiceUserSeed == "") { - return nil, fmt.Errorf("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") + return nil, fmt.Errorf("initializing cluster settings: %w", err) } - natsOpts := cfg.Distributed.NatsMessagingOptions("", "") - natsClient, err := messaging.New(cfg.Distributed.NatsURL, natsOpts...) + carrierRow, err := seedCarrier(context.Background(), carrierStore, settingsStore, cfg.Distributed.NatsURL, cfg.Distributed.InstanceID) if err != nil { - return nil, fmt.Errorf("connecting to NATS: %w", err) + return nil, err } - xlog.Info("Connected to NATS", "url", sanitize.URL(cfg.Distributed.NatsURL)) - // Ensure NATS is closed if any subsequent initialization step fails. + // Whatever fails from here on must give back what was built. success := false + var startSet *carrier.Set defer func() { - if !success { - natsClient.Close() + if !success && startSet != nil && startSet.Close != nil { + startSet.Close() } }() @@ -318,43 +329,92 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("File manager initialized", "cacheDir", cacheDir) - // Build the NATS carrier set and put it behind the holders. Everything - // below takes a holder, not the connection. - natsSet, err := carrier.NewNATSSet(carrier.NATSOptions{ - Client: natsClient, - Epoch: carrierRow.Epoch, - Registry: registry, - InstallTimeout: cfg.Distributed.BackendInstallTimeoutOrDefault(), - UpgradeTimeout: cfg.Distributed.BackendUpgradeTimeoutOrDefault(), - Token: cfg.Distributed.RegistrationToken, - S3Staging: cfg.Distributed.StorageURL != "", - FileManager: fileMgr, - HTTPAddrFor: func(nodeID string) (string, error) { - node, err := registry.Get(context.Background(), nodeID) - if err != nil { - return "", err - } - if node.HTTPAddress == "" { - return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) - } - return node.HTTPAddress, nil - }, - }) + // The pieces of the tunnel carrier that do not depend on a set. They exist on + // every replica, whatever the carrier in use: a change of carrier needs them, + // and the routes of the tunnel are registered on every replica. + jobStore, err := jobs.NewJobStore(authDB) if err != nil { - return nil, fmt.Errorf("building the %s carrier: %w", cluster.CarrierNATS, err) + return nil, fmt.Errorf("initializing job store: %w", err) + } + // The claim queue is the queue of the tunnel carrier. Its table exists on + // every carrier, so that a change of carrier finds it. + if err := jobs.MigrateClaims(context.Background(), authDB); err != nil { + return nil, err } + xlog.Info("Distributed job store initialized") + + // A replica proves who it is to its peers with a credential that it mints + // itself. Only the hash is published, in the same row as the address that + // the peers dial. + peerCred := cluster.NewPeerCredential() + advertised := peerAddress(cfg) + membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID, advertised, peerCred) + if err != nil { + return nil, err + } + defer func() { + if !success { + membership.Stop() + } + }() + clusterReg := cluster.NewRegistry(authDB) + tunnels := tunnel.NewRegistry(clusterReg, cfg.Distributed.InstanceID) + // If a peer sweeps this replica while it stalls, the claims of the tunnels + // that it still holds are written again when it registers again. + membership.SetReclaimer(tunnels) + peerOpts, err := peerPoolOptions(cfg, advertised) + if err != nil { + return nil, err + } + peerPool := tunnel.NewPeerPool(cfg.Distributed.InstanceID, peerCred, clusterReg, peerOpts...) + peerSessions := tunnel.NewPeerSessions(tunnel.NewRelay(tunnels).Stream) + + // Build the set of the carrier that the row names, and put it behind the + // holders. Everything below takes a holder, not a carrier. A NATS server that + // is not up yet does not stop the start; a URL or a credential that cannot be + // used does. active := &atomic.Pointer[carrier.Set]{} - active.Store(natsSet) + runtime := &carrierRuntime{ + cfg: cfg, db: authDB, settings: settingsStore, registry: registry, fileMgr: fileMgr, + instanceID: cfg.Distributed.InstanceID, jobStore: jobStore, clusterReg: clusterReg, + tunnels: tunnels, peerPool: peerPool, active: active, + probeWake: make(chan struct{}, 1), + } + startSet, err = runtime.build(context.Background(), carrierRow, carrierRow.Active, false) + if err != nil { + return nil, fmt.Errorf("building the %s carrier that the cluster uses: %w", carrierRow.Active, err) + } + active.Store(startSet) broadcaster := carrier.NewBroadcaster(active) + runtime.bus = broadcaster + + // While two carriers are attached, a call to a worker goes to the carrier the + // worker is attached to. + workers := nodes.NewSwitchWorkers(registry, clusterReg, cluster.DefaultReconnectGrace, cfg.Distributed.StaleNodeThresholdOrDefault()) + window := carrier.NewWindow( + func(ctx context.Context, nodeID string) (carrier.Attachment, error) { + attached, err := workers.AttachedCarriers(ctx, nodeID) + if err != nil { + return carrier.Attachment{}, err + } + return attachmentOf(attached), nil + }, + workers.AgentsAttached, + ) workQueue := carrier.NewWorkQueue(active) clientFactory := carrier.NewClients(active) + clientFactory.UseWindow(window) fileStager := carrier.NewFiles(active) + fileStager.UseWindow(window) remoteUnloader := carrier.NewCommands(active) - workerHTTPDial := carrier.NewWorkerDialer(active) + remoteUnloader.UseWindow(window) + workerHTTPDial := carrier.NewRoutedWorkerDialer(active, window) + agentControl := carrier.NewAgents(active) + agentControl.UseWindow(window) if cfg.Distributed.StorageURL != "" { - xlog.Info("File stager initialized (S3+NATS)") + xlog.Info("File stager initialized (S3)", "carrier", carrierRow.Active) } else { - xlog.Info("File stager initialized (HTTP direct transfer)") + xlog.Info("File stager initialized (HTTP direct transfer)", "carrier", carrierRow.Active) } // Collect SmartRouter option values; the router itself is created after all @@ -376,18 +436,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade clientFactory, ) - // Initialize job store - jobStore, err := jobs.NewJobStore(authDB) - if err != nil { - return nil, fmt.Errorf("initializing job store: %w", err) - } - // The claim queue is the queue of the tunnel carrier. Its table exists on - // every carrier, so that a change of carrier finds it. - if err := jobs.MigrateClaims(context.Background(), authDB); err != nil { - return nil, err - } - xlog.Info("Distributed job store initialized") - // Initialize job dispatcher dispatcher := jobs.NewDispatcher(jobStore, workQueue, broadcaster, authDB, cfg.Distributed.InstanceID) @@ -602,16 +650,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Create ModelRouterAdapter to wire into ModelLoader modelAdapter := nodes.NewModelRouterAdapter(router) - // A replica proves who it is to its peers with a credential that it mints - // itself. Only the hash is published, in the same row as the address that - // the peers dial. - peerCred := cluster.NewPeerCredential() - advertised := peerAddress(cfg) - membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID, advertised, peerCred) - if err != nil { - return nil, err - } - clusterReg := cluster.NewRegistry(authDB) // While the tunnel is the active carrier, a node that heartbeats and holds no // tunnel is demoted once its departure is older than the grace. On NATS no // node holds a tunnel, and the monitor reads nothing. @@ -619,21 +657,58 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade set := active.Load() return set != nil && set.Name == cluster.CarrierTunnel }) - tunnels := tunnel.NewRegistry(clusterReg, cfg.Distributed.InstanceID) - // If a peer sweeps this replica while it stalls, the claims of the tunnels - // that it still holds are written again when it registers again. - membership.SetReclaimer(tunnels) - peerOpts, err := peerPoolOptions(cfg, advertised) + + // The protocol of a change of carrier, and what this replica does in one. The + // switch tells the other replicas of each move with a hint, so that they look + // at the row at once. The hint is a courtesy; the poll is what counts. + fallbackTimings := cluster.Timings{ + PrepareTimeout: cfg.Distributed.CarrierPrepareTimeout, + TransitionWindow: cfg.Distributed.CarrierTransitionWindow, + MaxDrain: cfg.Distributed.CarrierMaxDrain, + } + carrierSwitch, err := cluster.NewSwitch(cluster.SwitchOptions{ + Store: carrierStore, + Registry: clusterReg, + Workers: workers, + Work: workCounts{db: authDB}, + Timings: func() cluster.Timings { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + t, err := settingsStore.Timings(ctx, fallbackTimings) + if err != nil { + xlog.Warn("Could not read the waits of a change of carrier from the cluster settings; using the ones of this replica", "error", err) + return fallbackTimings + } + return t + }, + AvailabilityMaxAge: availabilityMaxAge, + OnChange: func(row cluster.CarrierRow) { + go func() { + if err := broadcaster.Publish(messaging.SubjectCarrierChanged, row); err != nil { + xlog.Debug("Could not send the hint that the cluster carrier row moved; the other replicas read it on their poll", "error", err) + } + }() + }, + }) + if err != nil { + return nil, err + } + swapper, err := carrier.NewSwapper(carrier.SwapperOptions{ + Cur: active, Bus: broadcaster, Window: window, + Rows: carrierStore, Ready: membership, Build: runtime.Build, + }) if err != nil { return nil, err } - peerPool := tunnel.NewPeerPool(cfg.Distributed.InstanceID, peerCred, clusterReg, peerOpts...) - peerSessions := tunnel.NewPeerSessions(tunnel.NewRelay(tunnels).Stream) success = true return &DistributedServices{ Membership: membership, Carriers: carrierStore, + Settings: settingsStore, + Switch: carrierSwitch, + Swapper: swapper, + Runtime: runtime, Tunnels: tunnels, Instances: clusterReg, PeerSessions: peerSessions, @@ -641,7 +716,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade WorkerDialer: tunnel.NewWorkerDialer(tunnels, peerPool), Broadcaster: broadcaster, WorkQueue: workQueue, - AgentControl: carrier.NewAgents(active), + AgentControl: agentControl, Store: store, Registry: registry, Router: router, @@ -661,10 +736,91 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade WorkerHTTPDial: workerHTTPDial, - active: active, + active: active, + window: window, + workers: workers, }, nil } +// StartCarrierSwitch starts what this replica does about changes of carrier: it +// follows the cluster carrier row, drives the protocol when it holds the +// leadership, and reports which carriers it could build. Call it once, after the +// services that use the carrier have started. +func (ds *DistributedServices) StartCarrierSwitch(ctx context.Context, db *gorm.DB) { + // The hints. Both are a courtesy that shortens a wait. + if _, err := ds.Broadcaster.Subscribe(messaging.SubjectCarrierChanged, func([]byte) { ds.Swapper.Wake() }); err != nil { + xlog.Warn("Could not listen for the hint that the cluster carrier changed; the poll still follows the row", "error", err) + } + if _, err := ds.Broadcaster.Subscribe(messaging.SubjectCarrierProbe, func([]byte) { ds.Runtime.wakeProbe() }); err != nil { + xlog.Warn("Could not listen for the request to report which carriers this replica can build", "error", err) + } + go ds.Swapper.Run(ctx) + go ds.Runtime.RunAvailability(ctx) + // One replica leads the protocol at a time. If it dies, the next tick on + // another replica goes on from the row. + go advisorylock.RunLeaderLoop(ctx, db, advisorylock.KeyCarrierSwitch, leaderTick, func() { + if err := ds.Switch.Drive(ctx); err != nil && ctx.Err() == nil { + xlog.Warn("The change of carrier could not advance", "error", err) + } + }) +} + +// leaderTick is how often the replica that leads makes the move that is due. +const leaderTick = 2 * time.Second + +// seedCarrier reads the cluster carrier row, and writes it when no replica has +// yet. It returns the row as it is after the call. +func seedCarrier(ctx context.Context, store *cluster.CarrierStore, settings *cluster.SettingsStore, flagURL, instanceID string) (cluster.CarrierRow, error) { + stored, _, err := settings.Get(ctx, cluster.SettingNATSURL) + if err != nil { + return cluster.CarrierRow{}, err + } + seedWith := cluster.CarrierTunnel + if flagURL != "" || stored != "" { + seedWith = cluster.CarrierNATS + } + row, created, err := store.Seed(ctx, seedWith, instanceID) + if err != nil { + return cluster.CarrierRow{}, fmt.Errorf("reading cluster carrier state: %w", err) + } + if created { + xlog.Info("Cluster carrier seeded", "carrier", row.Active, "epoch", row.Epoch) + } + if flagURL != "" { + // The URL of the flag is copied into the cluster settings when none is + // stored, so that every replica uses the same one. A stored URL wins. + made, err := settings.SetIfAbsent(ctx, cluster.SettingNATSURL, flagURL, instanceID) + if err != nil { + return cluster.CarrierRow{}, err + } + switch { + case made: + xlog.Info("The NATS URL of this replica is now a setting of the cluster", "url", sanitize.URL(flagURL)) + case stored != "" && stored != flagURL: + xlog.Warn("The NATS URL of this replica differs from the one stored for the cluster; the stored one is used", + "flag", sanitize.URL(flagURL), "cluster", sanitize.URL(stored)) + } + } + if row.Active == cluster.CarrierTunnel && flagURL != "" { + xlog.Info("The cluster runs on the tunnel carrier. The NATS URL of this replica is kept for a change back to NATS, and no NATS connection is opened") + } + return row, nil +} + +// attachmentOf turns the carriers a worker is attached to into the routing form. +func attachmentOf(cs []cluster.Carrier) carrier.Attachment { + var a carrier.Attachment + for _, c := range cs { + switch c { + case cluster.CarrierNATS: + a.NATS = true + case cluster.CarrierTunnel: + a.Tunnel = true + } + } + return a +} + // startMembership creates the cluster tables, then records this replica in the // instances table and keeps its row fresh until the Membership is stopped. // Every replica does this on every carrier, because the change of carrier needs diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go index c03fdde573af..8a2eb7820e0d 100644 --- a/core/application/distributed_carrier_test.go +++ b/core/application/distributed_carrier_test.go @@ -13,6 +13,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" "github.com/mudler/LocalAI/core/services/tunnel" "github.com/testcontainers/testcontainers-go" @@ -67,28 +68,107 @@ var _ = Describe("distributed startup and the cluster carrier", func() { Expect(row.ChangedBy).To(Equal("replica-a")) }) - It("refuses to start on a carrier this build does not run, and does not try NATS", func() { - store, err := cluster.NewCarrierStore(db) + It("starts a deployment with only PostgreSQL on the tunnel carrier, and opens no NATS connection", func() { + pg, dsn := testutil.SetupTestDBWithDSN() + cfg.Auth.DatabaseURL = dsn + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, pg, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + store, err := cluster.NewCarrierStore(pg) + Expect(err).ToNot(HaveOccurred()) + row, err := store.Get(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Active).To(Equal(cluster.CarrierTunnel)) + Expect(row.Epoch).To(Equal(int64(1))) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierTunnel)) + + // Fan-out goes through the database, and one LISTEN session exists for it. + var sessions int64 + Expect(pg.Raw("SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'").Scan(&sessions).Error).To(Succeed()) + Expect(sessions).To(Equal(int64(1))) + got := make(chan []byte, 1) + _, err = svc.Broadcaster.Subscribe("jobs.carrier-spec", func(b []byte) { got <- b }) + Expect(err).ToNot(HaveOccurred()) + Expect(svc.Broadcaster.Publish("jobs.carrier-spec", "hello")).To(Succeed()) + Eventually(got, "10s").Should(Receive(Equal([]byte(`"hello"`)))) + }) + + It("starts on the tunnel carrier that the row names, whatever the flags of this replica say, and does not try NATS", func() { + pg, dsn := testutil.SetupTestDBWithDSN() + cfg.Auth.DatabaseURL = dsn + store, err := cluster.NewCarrierStore(pg) Expect(err).ToNot(HaveOccurred()) _, _, err = store.Seed(context.Background(), cluster.CarrierTunnel, "replica-b") Expect(err).ToNot(HaveOccurred()) - svc, err := initDistributed(cfg, db, nil, nil) - Expect(err).To(MatchError(ContainSubstring(`the cluster carrier is "tunnel"`))) - Expect(err.Error()).ToNot(ContainSubstring("connecting to NATS"), "no fallback to the carrier the flags name") - Expect(svc).To(BeNil()) - Expect(carrierRow().Active).To(Equal(cluster.CarrierTunnel), "the row is left as it was") + // The NATS URL of the flag points at nothing. It is kept for a change back to + // NATS, and it is not dialled. + svc, err := initDistributed(cfg, pg, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierTunnel)) + row, err := store.Get(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Active).To(Equal(cluster.CarrierTunnel), "the row is left as it was") + Expect(row.ChangedBy).To(Equal("replica-b")) + + url, ok, err := svc.Settings.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://127.0.0.1:1"), "the URL of the flag is stored for a change back to NATS") }) - It("keeps the existing validation of a missing NATS URL", func() { + It("claims queued work on the tunnel carrier once the switch has started", func() { + pg, dsn := testutil.SetupTestDBWithDSN() + cfg.Auth.DatabaseURL = dsn cfg.Distributed.NatsURL = "" - _, err := initDistributed(cfg, db, nil, nil) - Expect(err).To(MatchError(ContainSubstring("--nats-url"))) + svc, err := initDistributed(cfg, pg, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + svc.StartCarrierSwitch(cfg.Context, pg) + + Expect(svc.WorkQueue.Enqueue(context.Background(), messaging.WorkMCPCI, map[string]string{"job_id": "j1"})).To(Succeed()) + // No agent worker exists, so the claim goes back to the pool with a wait. That + // the attempt count moves shows that this replica took the row and drove it. + Eventually(func() int { + var attempts int + _ = pg.Raw("SELECT coalesce(max(attempts), 0) FROM work_claims").Scan(&attempts).Error + return attempts + }, "30s", "200ms").Should(BeNumerically(">=", 1)) + }) - store, err := cluster.NewCarrierStore(db) + It("seeds NATS for a deployment that has a NATS URL, and stores the URL for the whole cluster", func() { + // The server at that address is not there. A frontend may start before its + // broker, so the start goes on and the client keeps trying. + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierNATS)) + + row := carrierRow() + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + url, ok, err := store.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://127.0.0.1:1")) + }) + + It("uses the NATS URL stored for the cluster and not the one of the flag", func() { + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + Expect(store.Set(context.Background(), cluster.SettingNATSURL, "nats://[::1", "admin")).To(Succeed()) + cfg.Distributed.NatsURL = "nats://127.0.0.1:2" + + _, err = initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("connecting to NATS")), "the stored address, which cannot be parsed, was the one used") + url, _, err := store.Get(context.Background(), cluster.SettingNATSURL) Expect(err).ToNot(HaveOccurred()) - _, err = store.Get(context.Background()) - Expect(err).To(MatchError(cluster.ErrNotSeeded), "a replica that cannot start leaves no row behind") + Expect(url).To(Equal("nats://[::1"), "a flag does not overwrite a setting of the cluster") }) Context("with a NATS server", func() { diff --git a/core/application/startup.go b/core/application/startup.go index caaf78f7e83e..af3926456c59 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -339,6 +339,10 @@ func New(opts ...config.AppOption) (*Application, error) { if err := distSvc.Dispatcher.Start(options.Context); err != nil { return nil, fmt.Errorf("starting job dispatcher: %w", err) } + // Follow the cluster carrier row, lead the protocol of a change of carrier when + // this replica holds the leadership, and report which carriers it could build. + // It starts after the services that use the carrier. + distSvc.StartCarrierSwitch(options.Context, application.authDB) // Start ephemeral file cleanup storage.StartEphemeralCleanup(options.Context, distSvc.FileMgr, 0, 0) // Wire distributed backends into AgentJobService (before Start) diff --git a/core/services/advisorylock/keys.go b/core/services/advisorylock/keys.go index 76a1834eddf9..a87a6be2ab12 100644 --- a/core/services/advisorylock/keys.go +++ b/core/services/advisorylock/keys.go @@ -4,15 +4,18 @@ package advisorylock // These keys are global across the database — avoid collisions with // other applications sharing the same PostgreSQL instance. const ( - KeyCronScheduler int64 = 100 - KeyStaleNodeCleanup int64 = 101 - KeyGalleryDedup int64 = 102 - KeyAgentScheduler int64 = 103 - KeyHealthCheck int64 = 104 - KeySchemaMigrate int64 = 105 - KeyBackendUpgradeCheck int64 = 106 - KeyStateReconciler int64 = 107 + KeyCronScheduler int64 = 100 + KeyStaleNodeCleanup int64 = 101 + KeyGalleryDedup int64 = 102 + KeyAgentScheduler int64 = 103 + KeyHealthCheck int64 = 104 + KeySchemaMigrate int64 = 105 + KeyBackendUpgradeCheck int64 = 106 + KeyStateReconciler int64 = 107 // KeyFailoverProber elects the one frontend that probes failover // targets, decides chains and preloads warm targets. - KeyFailoverProber int64 = 108 + KeyFailoverProber int64 = 108 + // KeyCarrierSwitch elects the one frontend that makes the next move of a + // change of carrier. + KeyCarrierSwitch int64 = 109 ) diff --git a/core/services/carrier/nats.go b/core/services/carrier/nats.go index 768d898f671f..0213aed578a9 100644 --- a/core/services/carrier/nats.go +++ b/core/services/carrier/nats.go @@ -41,6 +41,12 @@ type NATSOptions struct { // HTTPAddrFor returns the address of a worker's file transfer server, for // the HTTP stager. HTTPAddrFor func(nodeID string) (string, error) + + // NodeHasAddress says whether a worker registered an address that can be + // dialled. With it, a worker that registered none (it holds a tunnel) is not + // dialled on a loopback address, which would name a port on the frontend and + // make a healthy backend look dead. Without it every address is dialled. + NodeHasAddress func(nodeID string) (bool, error) } // NewNATSSet builds the set of seam implementations that run over NATS for @@ -58,6 +64,11 @@ func NewNATSSet(o NATSOptions) (*Set, error) { files = nodes.NewHTTPFileStager(o.HTTPAddrFor, o.Token, dial) } + var clients nodes.BackendClientFactory = nodes.NewTokenClientFactory(o.Token) + if o.NodeHasAddress != nil { + clients = nodes.NewGuardedDirectClientFactory(o.Token, o.NodeHasAddress) + } + set := &Set{ Name: cluster.CarrierNATS, Epoch: o.Epoch, @@ -66,7 +77,7 @@ func NewNATSSet(o NATSOptions) (*Set, error) { WorkQueue: messaging.NewNATSWorkQueue(o.Client), Commands: nodes.NewRemoteUnloaderAdapter(o.Registry, o.Client, o.InstallTimeout, o.UpgradeTimeout), Files: files, - Clients: nodes.NewTokenClientFactory(o.Token), + Clients: clients, Dialer: dial, Agents: nodes.NewNATSAgentControl(o.Client), Close: o.Client.Close, diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 55c160112d5e..70b3f12184f8 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -27,6 +27,12 @@ const ( // says is what counts. const SubjectCarrierChanged = "state.carrier" +// SubjectCarrierProbe asks every replica to look again at which carriers it could +// build, and to write what it finds to the instances table. A dry run of a change +// of carrier sends it before it reads the answers. Like the other hint it is at +// most once, and a replica that does not hear it is reported as stale. +const SubjectCarrierProbe = "state.carrier.probe" + // SubjectClaimWake is the hint that a unit of work was put in the claim queue. // A consumer that hears it looks for work at once and does not wait for its next // poll. The hint is a broadcast and broadcasts are at-most-once, so it can be From fcd578c7b10b8568bb15fa8fb63d85e4831b9ec5 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:44:56 +0000 Subject: [PATCH 112/162] feat(http): add the admin API of the carrier GET /api/cluster/carrier reports the active carrier, the epoch, the state of a change, each live replica with its version and readiness, the workers that could not follow a change to the other carrier and why, the work in flight, the waits of a change and the live replica list that the admin confirms before a change. POST /api/cluster/carrier takes a target with dry_run and force, or abort. A dry run asks every replica to look at what it can build, and answers from what they wrote, and a request runs the same preflight and starts the change. The answer is 202 because the replicas carry out the change and one of them leads it. A blocked request is 422 with the blockers and the report, a change already under way is 409. GET and PUT /api/cluster/settings store the NATS address that frontends use, the one workers are told to use, and the waits of a change. Saving a NATS address makes NATS available and checks that this replica reaches it with its own credentials. It does not switch, and the switch stays an explicit request. A password in an address is masked when it is read back. No credential is stored. The routes sit under /api/cluster/, which is not a public prefix. Only the exact paths of the connect and peer routes skip the global authentication, and these four answer to an admin, so an anonymous caller gets 401 on each of them. They are registered on every frontend, and a frontend that is not distributed answers 503. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/carrier_runtime.go | 42 ++- core/http/app.go | 5 + .../http/endpoints/localai/cluster_carrier.go | 322 ++++++++++++++++++ .../endpoints/localai/cluster_carrier_test.go | 296 ++++++++++++++++ core/http/route_coverage_test.go | 27 ++ core/http/routes/cluster.go | 44 +++ 6 files changed, 727 insertions(+), 9 deletions(-) create mode 100644 core/http/endpoints/localai/cluster_carrier.go create mode 100644 core/http/endpoints/localai/cluster_carrier_test.go diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go index 02de689c0af0..a5fb8703155d 100644 --- a/core/application/carrier_runtime.go +++ b/core/application/carrier_runtime.go @@ -106,23 +106,30 @@ func (rt *carrierRuntime) natsClient(ctx context.Context, wait time.Duration) (* if url == "" { return nil, "", errors.New("no NATS URL is configured: store one in the cluster settings (nats.url) or start the frontend with --nats-url") } + client, err := rt.natsClientAt(ctx, url, wait) + return client, url, err +} + +// natsClientAt connects to the NATS server at url with the credentials of this +// replica. With a wait it returns only once the server has answered. +func (rt *carrierRuntime) natsClientAt(ctx context.Context, url string, wait time.Duration) (*messaging.Client, error) { d := rt.cfg.Distributed auth := d.NatsAuthConfig() if auth.RequireAuth && (auth.ServiceUserJWT == "" || auth.ServiceUserSeed == "") { - return nil, "", errors.New("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") + return nil, errors.New("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") } // A URL or a credential the client cannot use stops here. A server that is // not up does not, unless the caller waits for it: the client keeps trying. client, err := messaging.New(url, d.NatsMessagingOptions("", "")...) if err != nil { - return nil, "", fmt.Errorf("connecting to NATS: %w", err) + return nil, fmt.Errorf("connecting to NATS: %w", err) } if wait > 0 { deadline := time.Now().Add(wait) for !client.IsConnected() { if time.Now().After(deadline) || ctx.Err() != nil { client.Close() - return nil, "", fmt.Errorf("the NATS server at %s did not answer within %s: check the address and the credentials", sanitize.URL(url), wait) + return nil, fmt.Errorf("the NATS server at %s did not answer within %s: check the address and the credentials", sanitize.URL(url), wait) } select { case <-ctx.Done(): @@ -130,7 +137,18 @@ func (rt *carrierRuntime) natsClient(ctx context.Context, wait time.Duration) (* } } } - return client, url, nil + return client, nil +} + +// CheckNATS says whether this replica reaches the NATS server at url with its own +// credentials. It opens a connection and closes it. +func (rt *carrierRuntime) CheckNATS(ctx context.Context, url string) error { + client, err := rt.natsClientAt(ctx, url, natsProbeWait) + if err != nil { + return err + } + client.Close() + return nil } func (rt *carrierRuntime) buildNATS(ctx context.Context, row cluster.CarrierRow, wait bool) (*carrier.Set, error) { @@ -315,11 +333,9 @@ func (rt *carrierRuntime) wakeProbe() { } } -// ProbeReplicas asks every replica to look at what it can build, and waits until -// the live replicas have answered or probeWait has passed. A replica that did not -// answer shows up as stale in the preflight. -func (rt *carrierRuntime) ProbeReplicas(ctx context.Context) { - asked := time.Now() +// AskReplicas asks every replica, this one included, to look at what it can +// build, and does not wait for the answer. +func (rt *carrierRuntime) AskReplicas(ctx context.Context) { rt.wakeProbe() pubCtx, cancel := context.WithTimeout(ctx, hintPublishTimeout) defer cancel() @@ -334,6 +350,14 @@ func (rt *carrierRuntime) ProbeReplicas(ctx context.Context) { case <-done: case <-pubCtx.Done(): } +} + +// ProbeReplicas asks every replica to look at what it can build, and waits until +// the live replicas have answered or probeWait has passed. A replica that did not +// answer shows up as stale in the preflight. +func (rt *carrierRuntime) ProbeReplicas(ctx context.Context) { + asked := time.Now() + rt.AskReplicas(ctx) deadline := asked.Add(probeWait) for time.Now().Before(deadline) && ctx.Err() == nil { diff --git a/core/http/app.go b/core/http/app.go index cf65e6b8427d..fe9ac2091ba4 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -587,6 +587,11 @@ func API(application *application.Application) (*echo.Echo, error) { onPeer = d.PeerSessions.Accept } routes.RegisterClusterRoutes(e, registry, tunnels, instances, onPeer) + var clusterAdmin routes.ClusterAdmin + if d := application.Distributed(); d != nil { + clusterAdmin = routes.ClusterAdmin{Switch: d.Switch, Settings: d.Settings, Prober: d.Runtime, NATS: d.Runtime} + } + routes.RegisterClusterAdminRoutes(e, adminMiddleware, clusterAdmin) routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDial, registerOpts...) // Distributed SSE routes (job progress + agent events via NATS) diff --git a/core/http/endpoints/localai/cluster_carrier.go b/core/http/endpoints/localai/cluster_carrier.go new file mode 100644 index 000000000000..49b2809dafaf --- /dev/null +++ b/core/http/endpoints/localai/cluster_carrier.go @@ -0,0 +1,322 @@ +package localai + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/url" + "slices" + "time" + + "github.com/labstack/echo/v4" + + "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/pkg/sanitize" + "github.com/mudler/xlog" +) + +// CarrierSwitchAdmin is the protocol of a change of carrier, as the admin API +// uses it. *cluster.Switch is one. +type CarrierSwitchAdmin interface { + Status(ctx context.Context) (cluster.Report, error) + Preflight(ctx context.Context, target cluster.Carrier) (cluster.Report, error) + Request(ctx context.Context, req cluster.Request) (cluster.CarrierRow, cluster.Report, error) + Abort(ctx context.Context, by string) (cluster.CarrierRow, error) +} + +// ReplicaProber asks the replicas to look at which carriers they could build. +type ReplicaProber interface { + // ProbeReplicas asks and waits for the answers, or for a time limit. + ProbeReplicas(ctx context.Context) + // AskReplicas asks and does not wait. + AskReplicas(ctx context.Context) +} + +// NATSChecker tells whether this replica can reach a NATS server with its own +// credentials. +type NATSChecker interface { + CheckNATS(ctx context.Context, url string) error +} + +// ClusterSettings is the store of the settings of the cluster. +// *cluster.SettingsStore is one. +type ClusterSettings interface { + Get(ctx context.Context, key string) (string, bool, error) + Set(ctx context.Context, key, value, by string) error + Clear(ctx context.Context, key string) error + All(ctx context.Context) (map[string]string, error) +} + +func actor(c echo.Context) string { + if u := auth.GetUser(c); u != nil { + if u.Email != "" { + return u.Email + } + return u.ID + } + return "admin" +} + +func carrierError(c echo.Context, status int, msg string, extra map[string]any) error { + body := map[string]any{"error": msg} + for k, v := range extra { + body[k] = v + } + return c.JSON(status, body) +} + +// settingsView is the settings as an admin reads them. A password in the address +// of a server is masked. +func settingsView(all map[string]string) map[string]any { + view := map[string]any{ + "nats_url": "", + "nats_worker_url": "", + "prepare_timeout": all[cluster.SettingPrepareTimeout], + "transition_window": all[cluster.SettingTransitionWindow], + "max_drain": all[cluster.SettingMaxDrain], + } + if v := all[cluster.SettingNATSURL]; v != "" { + view["nats_url"] = sanitize.URL(v) + } + if v := all[cluster.SettingNATSWorkerURL]; v != "" { + view["nats_worker_url"] = sanitize.URL(v) + } + return view +} + +// carrierView is the report with the settings next to it. +type carrierView struct { + cluster.Report + Settings map[string]any `json:"settings"` +} + +// GetCarrierEndpoint reports the state of the carrier of the cluster: the active +// carrier, the epoch, the state of a change, each replica and whether it is ready, +// the workers that could not follow a change to the other carrier and why, the +// work in flight, and the list of live replicas for the admin to confirm before a +// change. +func GetCarrierEndpoint(sw CarrierSwitchAdmin, settings ClusterSettings) echo.HandlerFunc { + return func(c echo.Context) error { + ctx := c.Request().Context() + status, err := sw.Status(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + // The report is for a change to the other carrier, so the workers that could + // not follow it are listed with their reasons. + other := cluster.CarrierTunnel + if status.Active == cluster.CarrierTunnel { + other = cluster.CarrierNATS + } + report, err := sw.Preflight(ctx, other) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + all, err := settings.All(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + return c.JSON(http.StatusOK, carrierView{Report: report, Settings: settingsView(all)}) + } +} + +type switchRequest struct { + Target string `json:"target"` + DryRun bool `json:"dry_run"` + Force bool `json:"force"` + Abort bool `json:"abort"` +} + +// SwitchCarrierEndpoint changes the carrier of the cluster. It takes +// {"target": "nats"|"tunnel", "dry_run": bool, "force": bool}, or {"abort": true} +// to abort a change that is being prepared. +// +// A dry run asks every replica to look at what it can build, runs the preflight, +// and changes nothing. A request runs the same preflight and starts the change +// when nothing blocks. The change itself is carried out by the replicas and led by +// one of them, so the answer is 202: it says what was started and the epoch to +// watch. +func SwitchCarrierEndpoint(sw CarrierSwitchAdmin, prober ReplicaProber) echo.HandlerFunc { + return func(c echo.Context) error { + var req switchRequest + if err := c.Bind(&req); err != nil { + return carrierError(c, http.StatusBadRequest, "invalid request body", nil) + } + ctx := c.Request().Context() + by := actor(c) + + if req.Abort { + row, err := sw.Abort(ctx, by) + switch { + case errors.Is(err, cluster.ErrNotAbortable): + return carrierError(c, http.StatusConflict, err.Error(), nil) + case err != nil: + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + xlog.Warn("Change of carrier aborted by an admin", "by", by, "epoch", row.Epoch) + return c.JSON(http.StatusOK, row) + } + + target := cluster.Carrier(req.Target) + if target != cluster.CarrierNATS && target != cluster.CarrierTunnel { + return carrierError(c, http.StatusBadRequest, fmt.Sprintf("target must be %q or %q", cluster.CarrierNATS, cluster.CarrierTunnel), nil) + } + // The replicas write what they can build when they are asked, so the answer + // of the preflight is about now and not about the last time they looked. + prober.ProbeReplicas(ctx) + + if req.DryRun { + report, err := sw.Preflight(ctx, target) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + return c.JSON(http.StatusOK, report) + } + + row, report, err := sw.Request(ctx, cluster.Request{Target: target, By: by, Force: req.Force}) + var blocked *cluster.BlockedError + switch { + case errors.As(err, &blocked): + return carrierError(c, http.StatusUnprocessableEntity, err.Error(), map[string]any{"blockers": blocked.Report.Blockers, "report": blocked.Report}) + case errors.Is(err, cluster.ErrBusy): + return carrierError(c, http.StatusConflict, err.Error(), nil) + case err != nil: + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + xlog.Info("Change of carrier started by an admin", "by", by, "target", target, "force", req.Force, "epoch", row.Epoch) + return c.JSON(http.StatusAccepted, map[string]any{ + "state": row.State, "active": row.Active, "target": row.Target, "epoch": row.Epoch, + "force": row.Force, "warnings": report.Warnings, "row": row, + }) + } +} + +// GetClusterSettingsEndpoint reads the settings of the cluster. +func GetClusterSettingsEndpoint(settings ClusterSettings) echo.HandlerFunc { + return func(c echo.Context) error { + all, err := settings.All(c.Request().Context()) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + return c.JSON(http.StatusOK, settingsView(all)) + } +} + +type settingsRequest struct { + NATSURL *string `json:"nats_url"` + NATSWorkerURL *string `json:"nats_worker_url"` + PrepareTimeout *string `json:"prepare_timeout"` + TransitionWindow *string `json:"transition_window"` + MaxDrain *string `json:"max_drain"` +} + +func validNATSURL(raw string) error { + u, err := url.Parse(raw) + if err != nil { + return fmt.Errorf("not a URL: %w", err) + } + if !slices.Contains([]string{"nats", "tls", "ws", "wss"}, u.Scheme) { + return fmt.Errorf("the scheme must be nats, tls, ws or wss, not %q", u.Scheme) + } + if u.Host == "" { + return errors.New("the address has no host") + } + return nil +} + +// natsCheckTimeout bounds the reachability check of a saved URL. +const natsCheckTimeout = 8 * time.Second + +// PutClusterSettingsEndpoint stores settings of the cluster. Saving a NATS URL +// makes NATS available and checks that this replica reaches it; it does not +// switch the cluster to NATS. The change of carrier is a separate request with a +// dry run. Every replica checks the stored URL on its next report, and a dry run +// asks them to. +func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, prober ReplicaProber) echo.HandlerFunc { + return func(c echo.Context) error { + var req settingsRequest + if err := c.Bind(&req); err != nil { + return carrierError(c, http.StatusBadRequest, "invalid request body", nil) + } + ctx := c.Request().Context() + by := actor(c) + + type change struct { + key string + val *string + url bool + } + changes := []change{ + {cluster.SettingNATSURL, req.NATSURL, true}, + {cluster.SettingNATSWorkerURL, req.NATSWorkerURL, true}, + {cluster.SettingPrepareTimeout, req.PrepareTimeout, false}, + {cluster.SettingTransitionWindow, req.TransitionWindow, false}, + {cluster.SettingMaxDrain, req.MaxDrain, false}, + } + // Validate everything before storing anything. + for _, ch := range changes { + if ch.val == nil || *ch.val == "" { + continue + } + if ch.url { + if err := validNATSURL(*ch.val); err != nil { + return carrierError(c, http.StatusBadRequest, fmt.Sprintf("%s: %v", ch.key, err), nil) + } + } else if _, err := time.ParseDuration(*ch.val); err != nil { + return carrierError(c, http.StatusBadRequest, fmt.Sprintf("%s: %v", ch.key, err), nil) + } + } + for _, ch := range changes { + if ch.val == nil { + continue + } + var err error + if *ch.val == "" { + err = settings.Clear(ctx, ch.key) + } else { + err = settings.Set(ctx, ch.key, *ch.val, by) + } + switch { + case errors.Is(err, cluster.ErrUnknownSetting): + return carrierError(c, http.StatusBadRequest, err.Error(), nil) + case err != nil: + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + } + xlog.Info("Cluster settings changed", "by", by) + + resp := map[string]any{"saved": true} + if req.NATSURL != nil && *req.NATSURL != "" { + checkCtx, cancel := context.WithTimeout(ctx, natsCheckTimeout) + err := checker.CheckNATS(checkCtx, *req.NATSURL) + cancel() + result := map[string]any{"reachable": err == nil} + if err != nil { + result["reason"] = err.Error() + xlog.Warn("The NATS URL that was saved is not reachable from this replica", "error", err) + } + resp["nats"] = result + } + if req.NATSURL != nil || req.NATSWorkerURL != nil { + // The other replicas look at the new address and write what they find. + prober.AskReplicas(ctx) + } + all, err := settings.All(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + resp["settings"] = settingsView(all) + return c.JSON(http.StatusOK, resp) + } +} + +// UnavailableCarrierEndpoint answers 503. The routes of the carrier are +// registered on every frontend, and a frontend that is not distributed has no +// cluster to change. +func UnavailableCarrierEndpoint() echo.HandlerFunc { + return func(c echo.Context) error { + return carrierError(c, http.StatusServiceUnavailable, "distributed mode is not enabled on this frontend", nil) + } +} diff --git a/core/http/endpoints/localai/cluster_carrier_test.go b/core/http/endpoints/localai/cluster_carrier_test.go new file mode 100644 index 000000000000..32018eabf24a --- /dev/null +++ b/core/http/endpoints/localai/cluster_carrier_test.go @@ -0,0 +1,296 @@ +package localai + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type fakeSwitch struct { + mu sync.Mutex + report cluster.Report + row cluster.CarrierRow + requests []cluster.Request + aborts []string + requestE error + abortE error + order *[]string +} + +func (f *fakeSwitch) note(s string) { + if f.order != nil { + *f.order = append(*f.order, s) + } +} + +func (f *fakeSwitch) Status(context.Context) (cluster.Report, error) { return f.report, nil } +func (f *fakeSwitch) Preflight(_ context.Context, target cluster.Carrier) (cluster.Report, error) { + f.note("preflight") + if target != cluster.CarrierNATS && target != cluster.CarrierTunnel { + return cluster.Report{}, cluster.ErrInvalidCarrier + } + r := f.report + r.Target = target + return r, nil +} + +func (f *fakeSwitch) Request(_ context.Context, req cluster.Request) (cluster.CarrierRow, cluster.Report, error) { + f.note("request") + f.mu.Lock() + f.requests = append(f.requests, req) + f.mu.Unlock() + return f.row, f.report, f.requestE +} + +func (f *fakeSwitch) Abort(_ context.Context, by string) (cluster.CarrierRow, error) { + f.mu.Lock() + f.aborts = append(f.aborts, by) + f.mu.Unlock() + return f.row, f.abortE +} + +type fakeProber struct{ order *[]string } + +func (f fakeProber) ProbeReplicas(context.Context) { + if f.order != nil { + *f.order = append(*f.order, "probe") + } +} + +func (f fakeProber) AskReplicas(context.Context) {} + +type fakeNATS struct { + err error + seen []string +} + +func (f *fakeNATS) CheckNATS(_ context.Context, url string) error { + f.seen = append(f.seen, url) + return f.err +} + +type memorySettings struct { + mu sync.Mutex + m map[string]string +} + +func (s *memorySettings) Get(_ context.Context, k string) (string, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.m[k] + return v, ok, nil +} + +func (s *memorySettings) Set(_ context.Context, k, v, _ string) error { + if k == "nats.password" { + return cluster.ErrUnknownSetting + } + s.mu.Lock() + defer s.mu.Unlock() + s.m[k] = v + return nil +} + +func (s *memorySettings) Clear(_ context.Context, k string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.m, k) + return nil +} + +func (s *memorySettings) All(context.Context) (map[string]string, error) { + s.mu.Lock() + defer s.mu.Unlock() + out := map[string]string{} + for k, v := range s.m { + out[k] = v + } + return out, nil +} + +var _ = Describe("The admin API of the carrier", func() { + var ( + sw *fakeSwitch + settings *memorySettings + nats *fakeNATS + order []string + ) + + call := func(h echo.HandlerFunc, method, body string) (int, map[string]any) { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(method, "/", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + Expect(h(e.NewContext(req, rec))).To(Succeed()) + var out map[string]any + if rec.Body.Len() > 0 { + Expect(json.Unmarshal(rec.Body.Bytes(), &out)).To(Succeed(), rec.Body.String()) + } + return rec.Code, out + } + + BeforeEach(func() { + order = nil + sw = &fakeSwitch{ + order: &order, + report: cluster.Report{Active: cluster.CarrierNATS, Epoch: 4, State: cluster.StateStable, OK: true, Replicas: []cluster.ReplicaStatus{{ID: "a", Version: "v1"}}}, + row: cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 5, State: cluster.StatePrepare, Target: cluster.CarrierTunnel}, + } + settings = &memorySettings{m: map[string]string{cluster.SettingNATSURL: "nats://user:secret@broker:4222"}} + nats = &fakeNATS{} + }) + + Describe("the status", func() { + It("reports the row, the replicas, the opposite carrier as the target, and a NATS URL without its password", func() { + code, body := call(GetCarrierEndpoint(sw, settings), http.MethodGet, "") + Expect(code).To(Equal(http.StatusOK)) + Expect(body["active"]).To(Equal("nats")) + Expect(body["epoch"]).To(BeEquivalentTo(4)) + Expect(body["target"]).To(Equal("tunnel"), "the report answers for a change to the other carrier") + Expect(body["replicas"]).To(HaveLen(1)) + shown := body["settings"].(map[string]any)["nats_url"].(string) + Expect(shown).To(ContainSubstring("broker:4222")) + Expect(shown).ToNot(ContainSubstring("secret")) + Expect(shown).ToNot(ContainSubstring("user")) + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("secret"), "the stored value is not changed") + }) + }) + + Describe("a request", func() { + It("asks the replicas to look, runs the preflight on a dry run, and changes nothing", func() { + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{order: &order}), http.MethodPost, `{"target":"tunnel","dry_run":true}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(body["target"]).To(Equal("tunnel")) + Expect(body["ok"]).To(BeTrue()) + Expect(order).To(Equal([]string{"probe", "preflight"})) + Expect(sw.requests).To(BeEmpty()) + }) + + It("starts the change, and passes force on", func() { + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{order: &order}), http.MethodPost, `{"target":"tunnel","force":true}`) + Expect(code).To(Equal(http.StatusAccepted)) + Expect(body["state"]).To(Equal("prepare")) + Expect(order).To(Equal([]string{"probe", "request"})) + Expect(sw.requests).To(HaveLen(1)) + Expect(sw.requests[0].Target).To(Equal(cluster.CarrierTunnel)) + Expect(sw.requests[0].Force).To(BeTrue()) + Expect(sw.requests[0].By).ToNot(BeEmpty(), "who asked is recorded") + }) + + It("answers 422 with the report when the preflight blocks", func() { + sw.requestE = &cluster.BlockedError{Report: cluster.Report{Blockers: []cluster.Blocker{{Kind: cluster.BlockerWorker, ID: "w1", Reason: "the worker predates carrier switching", Forceable: true}}}} + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusUnprocessableEntity)) + Expect(body["blockers"]).To(HaveLen(1)) + Expect(body["error"]).To(ContainSubstring("blocked")) + }) + + It("answers 409 when a change is already under way", func() { + sw.requestE = cluster.ErrBusy + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusConflict)) + Expect(body["error"]).To(ContainSubstring("under way")) + }) + + It("answers 400 for a target that is not a carrier, and for a body that is not JSON", func() { + code, _ := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"pigeon"}`) + Expect(code).To(Equal(http.StatusBadRequest)) + code, _ = call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `nope`) + Expect(code).To(Equal(http.StatusBadRequest)) + code, _ = call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"dry_run":true}`) + Expect(code).To(Equal(http.StatusBadRequest), "a request without a target") + }) + + It("aborts a change in prepare", func() { + sw.row = cluster.CarrierRow{Active: cluster.CarrierNATS, State: cluster.StateStable, Epoch: 6, Note: "aborted by admin"} + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"abort":true}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(body["state"]).To(Equal("stable")) + Expect(sw.aborts).To(HaveLen(1)) + Expect(sw.requests).To(BeEmpty()) + }) + + It("answers 409 to an abort when there is nothing to abort", func() { + sw.abortE = cluster.ErrNotAbortable + code, _ := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"abort":true}`) + Expect(code).To(Equal(http.StatusConflict)) + }) + + It("reports an error of the store as a server error", func() { + sw.requestE = errors.New("the database is down") + code, _ := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusInternalServerError)) + }) + }) + + Describe("the settings", func() { + It("stores a NATS URL, checks that this replica reaches it, and does not switch", func() { + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"nats_url":"nats://broker2:4222"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://broker2:4222")) + Expect(nats.seen).To(Equal([]string{"nats://broker2:4222"})) + Expect(body["nats"].(map[string]any)["reachable"]).To(BeTrue()) + Expect(sw.requests).To(BeEmpty()) + }) + + It("still stores a URL that this replica cannot reach, and says why", func() { + nats.err = errors.New("no answer within 3s") + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"nats_url":"nats://broker2:4222"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://broker2:4222")) + nb := body["nats"].(map[string]any) + Expect(nb["reachable"]).To(BeFalse()) + Expect(nb["reason"]).To(ContainSubstring("no answer")) + }) + + It("refuses a URL that is not a NATS address", func() { + for _, bad := range []string{`{"nats_url":"http://broker:4222"}`, `{"nats_url":"nats://"}`, `{"nats_url":"::"}`} { + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, bad) + Expect(code).To(Equal(http.StatusBadRequest), bad) + } + Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://user:secret@broker:4222")) + }) + + It("clears a URL with an empty string", func() { + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"nats_url":""}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m).ToNot(HaveKey(cluster.SettingNATSURL)) + }) + + It("stores the waits of a change, and refuses one that is not a positive duration", func() { + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"max_drain":"20m","prepare_timeout":"90s"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m[cluster.SettingMaxDrain]).To(Equal("20m")) + Expect(settings.m[cluster.SettingPrepareTimeout]).To(Equal("90s")) + code, _ = call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"max_drain":"soon"}`) + Expect(code).To(Equal(http.StatusBadRequest)) + }) + + It("reads the settings with the password of the URL masked", func() { + code, body := call(GetClusterSettingsEndpoint(settings), http.MethodGet, "") + Expect(code).To(Equal(http.StatusOK)) + Expect(body["nats_url"]).To(ContainSubstring("broker:4222")) + Expect(body["nats_url"]).ToNot(ContainSubstring("secret")) + }) + }) + + Describe("a frontend without distributed mode", func() { + It("answers 503 on every route", func() { + for _, h := range []echo.HandlerFunc{UnavailableCarrierEndpoint()} { + code, _ := call(h, http.MethodGet, "") + Expect(code).To(Equal(http.StatusServiceUnavailable)) + } + }) + }) +}) diff --git a/core/http/route_coverage_test.go b/core/http/route_coverage_test.go index 7505cb821f2f..6016be37cbbf 100644 --- a/core/http/route_coverage_test.go +++ b/core/http/route_coverage_test.go @@ -222,4 +222,31 @@ var _ = Describe("Route auth coverage", func() { app.ServeHTTP(rec, req) Expect(rec.Code).To(Equal(http.StatusUnauthorized)) }) + + // The admin routes of the carrier sit under /api/cluster/, next to the two + // routes that skip the global authentication. The prefix must not be public: + // an anonymous caller gets 401 on each of them, also on a frontend that is not + // distributed, and they exist whatever the mode. + It("registers the admin routes of the carrier and refuses an anonymous caller", func() { + want := map[string]bool{ + http.MethodGet + " /api/cluster/carrier": false, + http.MethodPost + " /api/cluster/carrier": false, + http.MethodGet + " /api/cluster/settings": false, + http.MethodPut + " /api/cluster/settings": false, + } + for _, r := range app.Routes() { + if _, ok := want[r.Method+" "+r.Path]; ok { + want[r.Method+" "+r.Path] = true + } + } + for route, found := range want { + Expect(found).To(BeTrue(), "%s is not registered", route) + method, path, _ := strings.Cut(route, " ") + req := httptest.NewRequest(method, path, strings.NewReader(`{"target":"tunnel"}`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusUnauthorized), route) + } + }) }) diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go index 34bc1484721e..8d1ebb7b2fe8 100644 --- a/core/http/routes/cluster.go +++ b/core/http/routes/cluster.go @@ -3,6 +3,7 @@ package routes import ( "github.com/labstack/echo/v4" clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/http/endpoints/localai" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/tunnel" @@ -25,3 +26,46 @@ func RegisterClusterRoutes(e *echo.Echo, registry *nodes.NodeRegistry, tunnels * e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(registry, tunnels)) e.GET(tunnel.PeerPath, clusterapi.PeerHandler(instances, onPeer)) } + +// Paths of the admin API of the carrier. They sit under /api/cluster/, which is +// not a public prefix: only the exact paths of the connect and peer routes skip +// the global authentication, and these answer to an admin only. +const ( + CarrierPath = "/api/cluster/carrier" + ClusterSettingsPath = "/api/cluster/settings" +) + +// ClusterAdmin is what the admin API of the carrier needs from the distributed +// services. A frontend that is not distributed passes nil values, and the routes +// answer 503. +type ClusterAdmin struct { + Switch localai.CarrierSwitchAdmin + Settings localai.ClusterSettings + Prober localai.ReplicaProber + NATS localai.NATSChecker +} + +// RegisterClusterAdminRoutes registers the admin API of the carrier: +// +// GET /api/cluster/carrier the state of the carrier, the replicas, the workers that could not follow +// POST /api/cluster/carrier a dry run, a change, or an abort +// GET /api/cluster/settings the settings of the cluster +// PUT /api/cluster/settings store a NATS address or the waits of a change +// +// The routes are registered on every frontend, as the connect route is, so that +// the route coverage test walks them and a frontend that is not distributed +// answers 503 and not 404. The admin middleware runs first. +func RegisterClusterAdminRoutes(e *echo.Echo, adminMw echo.MiddlewareFunc, admin ClusterAdmin) { + if admin.Switch == nil || admin.Settings == nil || admin.Prober == nil || admin.NATS == nil { + unavailable := localai.UnavailableCarrierEndpoint() + e.GET(CarrierPath, unavailable, adminMw) + e.POST(CarrierPath, unavailable, adminMw) + e.GET(ClusterSettingsPath, unavailable, adminMw) + e.PUT(ClusterSettingsPath, unavailable, adminMw) + return + } + e.GET(CarrierPath, localai.GetCarrierEndpoint(admin.Switch, admin.Settings), adminMw) + e.POST(CarrierPath, localai.SwitchCarrierEndpoint(admin.Switch, admin.Prober), adminMw) + e.GET(ClusterSettingsPath, localai.GetClusterSettingsEndpoint(admin.Settings), adminMw) + e.PUT(ClusterSettingsPath, localai.PutClusterSettingsEndpoint(admin.Settings, admin.NATS, admin.Prober), adminMw) +} From 749583957bfd2c3e3ebf794da990c7040e4154cb Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:58:08 +0000 Subject: [PATCH 113/162] test(carrier): change the carrier of a cluster with real sockets, a real broker and a real database The specs run frontend replicas the way the application builds them: a NATS set over a real server, a tunnel set over pgbus and the claim queue, the holders, the swapper, the window and a leader loop under the advisory lock. A worker answers on both carriers, and an agent worker consumes on each. From NATS to the tunnel, with an install, a held run, renewals of a load and chatter from both replicas in flight: the change settles in about half a second, every message is heard once by every replica, every job runs once on the carrier it was queued on, the install and the run end where they began, and no renewal fails. Back to NATS, the units that were queued on the tunnel run once on NATS and the LISTEN connection is closed. A replica that dies in prepare does not hold the change, and the one that starts again follows the row. A leader that dies between the commit and the settle is replaced by the other replica. An abort in prepare drops what was built for the change on every replica. A worker that cannot follow blocks the change, and a forced change leaves it reachable through the window and then unroutable, and reaps nothing. The claim loop and the handover of a tunnel set are built by NewTunnelSet from the options, so the specs and the application use the same code. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/carrier_runtime.go | 54 +- core/services/carrier/claims.go | 87 ++ .../carrier/switch_integration_test.go | 1008 +++++++++++++++++ core/services/carrier/tunnel.go | 9 + 4 files changed, 1110 insertions(+), 48 deletions(-) create mode 100644 core/services/carrier/claims.go create mode 100644 core/services/carrier/switch_integration_test.go diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go index a5fb8703155d..fe8649c5cfbb 100644 --- a/core/application/carrier_runtime.go +++ b/core/application/carrier_runtime.go @@ -2,7 +2,6 @@ package application import ( "context" - "encoding/json" "errors" "fmt" "sync/atomic" @@ -215,59 +214,18 @@ func (rt *carrierRuntime) buildTunnel(ctx context.Context, row cluster.CarrierRo Token: d.RegistrationToken, S3Staging: d.StorageURL != "", FileManager: rt.fileMgr, + // While the tunnel is the carrier in use, and while it drains, this replica + // claims queued work and drives it on agent workers. + Claims: &carrier.ClaimWork{ + DB: rt.db, Owner: rt.instanceID, Store: rt.jobStore, + Bus: func() messaging.Broadcaster { return rt.bus }, + }, }) if err != nil { fanout.Close() return nil, err } - // While the tunnel is the carrier in use, and while it drains, this replica - // claims queued work and drives it on agent workers. The loop dials through - // the dialer of this set and not through the holder: a run that started on the - // tunnel finishes there. - set.Start = func(ctx context.Context) (func(), error) { - loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{ - DB: rt.db, - Owner: rt.instanceID, - Picker: selector, - Control: nodes.NewControlClient(set.Dialer, d.RegistrationToken), - Broadcast: nodes.NewRebroadcaster(rt.bus), - Store: rt.jobStore, - Hints: rt.bus, - }) - if err != nil { - return nil, err - } - if err := loop.Start(ctx); err != nil { - return nil, err - } - xlog.Info("This replica claims queued work") - return loop.Stop, nil - } - - // When the tunnel is released, what is still queued moves to the queue of the - // carrier that took over. The rows are taken out of the table first and then - // published, so a job is lost, and not run twice, if this replica dies in - // between. A job that is lost stays running until the reaper fails it. - set.Handoff = func(ctx context.Context, next *carrier.Set) error { - rows, err := jobs.MigratePending(ctx, rt.db) - if err != nil { - return err - } - var failed error - moved := 0 - for _, r := range rows { - if err := next.WorkQueue.Enqueue(ctx, messaging.WorkKind(r.Kind), json.RawMessage(r.Payload)); err != nil { - failed = errors.Join(failed, fmt.Errorf("publishing the %s unit %s: %w", r.Kind, r.ID, err)) - continue - } - moved++ - } - if len(rows) > 0 { - xlog.Info("Queued work moved to the carrier that took over", "carrier", next.Name, "moved", moved, "lost", len(rows)-moved) - } - return failed - } return set, nil } diff --git a/core/services/carrier/claims.go b/core/services/carrier/claims.go new file mode 100644 index 000000000000..935cf7255069 --- /dev/null +++ b/core/services/carrier/claims.go @@ -0,0 +1,87 @@ +package carrier + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/mudler/xlog" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/jobs" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" +) + +// ClaimWork is what the tunnel carrier needs of a frontend replica besides its +// connections: the replica claims the queued work and drives it on agent +// workers, because an agent worker has no database and cannot claim. +type ClaimWork struct { + DB *gorm.DB + // Owner is the id of this replica in the instances table. + Owner string + // Store records the terminal state of the jobs that the work names. + Store jobs.TerminalWriter + // Bus returns the broadcaster holder. The hints of the queue and the events of + // the runs go through it, so that they follow the sets. It is a function + // because the holder is built after the first set, and it is read when the work + // starts. + Bus func() messaging.Broadcaster + // Interval is the poll of the claim loop. The default when zero. + Interval time.Duration +} + +// start returns the Start of a tunnel set. The loop dials through the dialer of +// that set and not through the holder: a run that started on the tunnel finishes +// there, whichever carrier is active by then. +func (c *ClaimWork) start(set func() *Set, selector *nodes.AgentSelector, token string) func(context.Context) (func(), error) { + return func(ctx context.Context) (func(), error) { + bus := c.Bus() + loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{ + DB: c.DB, + Owner: c.Owner, + Picker: selector, + Control: nodes.NewControlClient(set().Dialer, token), + Broadcast: nodes.NewRebroadcaster(bus), + Store: c.Store, + Hints: bus, + Interval: c.Interval, + }) + if err != nil { + return nil, err + } + if err := loop.Start(ctx); err != nil { + return nil, err + } + xlog.Info("This replica claims queued work", "owner", c.Owner) + return loop.Stop, nil + } +} + +// handoff returns the Handoff of a tunnel set: what is still queued moves to the +// queue of the carrier that took over. The rows are taken out of the table first +// and then published, so a job is lost, and not run twice, if this replica dies in +// between. A job that is lost stays running until the reaper fails it. +func (c *ClaimWork) handoff() func(context.Context, *Set) error { + return func(ctx context.Context, next *Set) error { + rows, err := jobs.MigratePending(ctx, c.DB) + if err != nil { + return err + } + var failed error + moved := 0 + for _, r := range rows { + if err := next.WorkQueue.Enqueue(ctx, messaging.WorkKind(r.Kind), json.RawMessage(r.Payload)); err != nil { + failed = errors.Join(failed, fmt.Errorf("publishing the %s unit %s: %w", r.Kind, r.ID, err)) + continue + } + moved++ + } + if len(rows) > 0 { + xlog.Info("Queued work moved to the carrier that took over", "carrier", next.Name, "moved", moved, "lost", len(rows)-moved) + } + return failed + } +} diff --git a/core/services/carrier/switch_integration_test.go b/core/services/carrier/switch_integration_test.go new file mode 100644 index 000000000000..50843d18fafc --- /dev/null +++ b/core/services/carrier/switch_integration_test.go @@ -0,0 +1,1008 @@ +package carrier_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http/httptest" + "os" + "runtime" + "slices" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/labstack/echo/v4" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" + ggrpc "google.golang.org/grpc" + "gorm.io/gorm" + + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/advisorylock" + "github.com/mudler/LocalAI/core/services/agentworker" + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/jobs" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/worker" + "github.com/mudler/LocalAI/core/services/workerctl" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var ( + switchNATSOnce sync.Once + switchNATSURL string + switchNATSCtr testcontainers.Container + switchNATSErr error +) + +// startSwitchNATS starts one NATS server for the process. +func startSwitchNATS() (string, error) { + switchNATSOnce.Do(func() { + ctx := context.Background() + var ctr testcontainers.Container + ctr, switchNATSErr = testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: testcontainers.ContainerRequest{ + Image: "nats:2.10-alpine", + ExposedPorts: []string{"4222/tcp"}, + WaitingFor: wait.ForListeningPort("4222/tcp"), + }, + Started: true, + }) + if switchNATSErr != nil { + return + } + switchNATSCtr = ctr + host, err := ctr.Host(ctx) + if err != nil { + switchNATSErr = err + return + } + port, err := ctr.MappedPort(ctx, "4222/tcp") + if err != nil { + switchNATSErr = err + return + } + switchNATSURL = fmt.Sprintf("nats://%s:%s", host, port.Port()) + }) + return switchNATSURL, switchNATSErr +} + +var _ = AfterSuite(func() { + if switchNATSCtr != nil { + _ = switchNATSCtr.Terminate(context.Background()) + } +}) + +// runLog records every run of a job and can hold a run until a spec lets it go. +type runLog struct { + mu sync.Mutex + runs map[string][]string + gates map[string]chan struct{} +} + +func newRunLog() *runLog { + return &runLog{runs: map[string][]string{}, gates: map[string]chan struct{}{}} +} + +func (r *runLog) hold(jobID string) chan struct{} { + r.mu.Lock() + defer r.mu.Unlock() + g := make(chan struct{}) + r.gates[jobID] = g + return g +} + +func (r *runLog) ranOn(jobID string) []string { + r.mu.Lock() + defer r.mu.Unlock() + return slices.Clone(r.runs[jobID]) +} + +func (r *runLog) total() int { + r.mu.Lock() + defer r.mu.Unlock() + n := 0 + for _, v := range r.runs { + n += len(v) + } + return n +} + +func (r *runLog) handler(label string) messaging.WorkHandler { + return func(ctx context.Context, payload []byte, events messaging.Publisher) error { + var evt jobs.JobEvent + if err := json.Unmarshal(payload, &evt); err != nil { + return err + } + r.mu.Lock() + r.runs[evt.JobID] = append(r.runs[evt.JobID], label) + gate := r.gates[evt.JobID] + r.mu.Unlock() + if gate != nil { + select { + case <-gate: + case <-ctx.Done(): + return ctx.Err() + } + } + jobs.PublishJobResult(events, evt.JobID, "completed", label, "") + return nil + } +} + +// tally counts what a subscriber heard, by message. +type tally struct { + mu sync.Mutex + by map[string]int +} + +func newTally() *tally { return &tally{by: map[string]int{}} } + +func (t *tally) handler(raw []byte) { + var msg string + if err := json.Unmarshal(raw, &msg); err != nil { + msg = string(raw) + } + t.mu.Lock() + defer t.mu.Unlock() + t.by[msg]++ +} + +func (t *tally) get(msg string) int { + t.mu.Lock() + defer t.mu.Unlock() + return t.by[msg] +} + +func (t *tally) copy() map[string]int { + t.mu.Lock() + defer t.mu.Unlock() + out := make(map[string]int, len(t.by)) + for k, v := range t.by { + out[k] = v + } + return out +} + +const chatterSubject = "state.switch-spec" + +// rigWork counts the jobs that run, which is all the specs read of the work in +// flight. +type rigWork struct{ db *gorm.DB } + +func (w rigWork) InFlight(ctx context.Context) (cluster.InFlight, error) { + var n int64 + err := w.db.WithContext(ctx).Model(&jobs.JobRecord{}).Where("status = ?", "running").Count(&n).Error + return cluster.InFlight{Jobs: int(n)}, err +} + +// switchRig is a cluster in one process: a database, a NATS server, a worker that +// can be reached over both carriers, an agent worker on each carrier, and +// frontend replicas that are built the way the application builds them. +type switchRig struct { + ctx context.Context + db *gorm.DB + dsn string + natsURL string + store *cluster.CarrierStore + clusterR *cluster.Registry + nodeReg *nodes.NodeRegistry + jobStore *jobs.JobStore + runs *runLog + + timings atomic.Pointer[cluster.Timings] + + // The worker. + nodeID string + backend *fakeBackend + control *controlPlane + grpcAddr string + httpAddr string + + // What the worker answered on each carrier. + natsRenewals, natsInstalls atomic.Int32 + installGate chan struct{} + installDone chan error + + workerNATS *messaging.Client + + // The agent worker. + agentID string + + replicas []*switchReplica + + // onMove, when set, is called by the replica whose switch made a move of the + // row. A spec uses it to kill the leader between two moves. + onMove atomic.Pointer[func(r *switchReplica, row cluster.CarrierRow)] +} + +type switchReplica struct { + rig *switchRig + id string + url string + tunnels *tunnel.Registry + sessions *tunnel.PeerSessions + pool *tunnel.PeerPool + cur atomic.Pointer[carrier.Set] + bus *carrier.Broadcaster + commands *carrier.Commands + queue *carrier.WorkQueue + window *carrier.Window + swapper *carrier.Swapper + sw *cluster.Switch + member *cluster.Membership + cancel context.CancelFunc + chatter *tally + dispatch *jobs.Dispatcher + mu sync.Mutex + builds map[cluster.Carrier]int + natsConns []*messaging.Client + failBuild map[cluster.Carrier]error + holdBuild map[cluster.Carrier]chan struct{} +} + +func (r *switchReplica) active() cluster.Carrier { return r.cur.Load().Name } + +func (r *switchReplica) buildCount(c cluster.Carrier) int { + r.mu.Lock() + defer r.mu.Unlock() + return r.builds[c] +} + +func (r *switchReplica) build(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*carrier.Set, error) { + rig := r.rig + r.mu.Lock() + err, hold := r.failBuild[target], r.holdBuild[target] + r.mu.Unlock() + if hold != nil { + select { + case <-hold: + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if err != nil { + return nil, err + } + r.mu.Lock() + r.builds[target]++ + r.mu.Unlock() + switch target { + case cluster.CarrierNATS: + client, err := messaging.New(rig.natsURL) + if err != nil { + return nil, err + } + for deadline := time.Now().Add(10 * time.Second); !client.IsConnected(); { + if time.Now().After(deadline) { + client.Close() + return nil, errors.New("the NATS server did not answer") + } + time.Sleep(20 * time.Millisecond) + } + r.mu.Lock() + r.natsConns = append(r.natsConns, client) + r.mu.Unlock() + return carrier.NewNATSSet(carrier.NATSOptions{ + Client: client, Epoch: row.Epoch, Registry: rig.nodeReg, + InstallTimeout: 30 * time.Second, UpgradeTimeout: 30 * time.Second, Token: registrationToken, + NodeHasAddress: func(nodeID string) (bool, error) { + n, err := rig.nodeReg.Get(context.Background(), nodeID) + if err != nil { + return false, err + } + return n.Address != "", nil + }, + }) + case cluster.CarrierTunnel: + fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: rig.db, DSN: rig.dsn}) + if err != nil { + return nil, err + } + return carrier.NewTunnelSet(carrier.TunnelOptions{ + Epoch: row.Epoch, + Dialer: tunnel.NewWorkerDialer(r.tunnels, r.pool), + Fanout: fan, + WorkQueue: jobs.NewClaimQueue(rig.db, fan.Broadcaster), + AgentSelector: nodes.NewAgentSelector(rig.nodeReg, rig.clusterR, r.id), + Registry: rig.nodeReg, + InstallTimeout: 30 * time.Second, UpgradeTimeout: 30 * time.Second, Token: registrationToken, + Claims: &carrier.ClaimWork{ + DB: rig.db, Owner: r.id, Store: rig.jobStore, + Bus: func() messaging.Broadcaster { return r.bus }, + Interval: 100 * time.Millisecond, + }, + }) + } + return nil, fmt.Errorf("unknown carrier %q", target) +} + +// kill stops everything the process of the replica was doing, without telling +// anyone: it does not leave the instances table and does not close its carrier. +func (r *switchReplica) kill() { + r.cancel() + Expect(r.rig.db.Exec(`UPDATE instances SET last_seen = now() - interval '1 hour' WHERE id = ?`, r.id).Error).To(Succeed()) +} + +func (rig *switchRig) startReplica(id string) *switchReplica { + GinkgoHelper() + ctx, cancel := context.WithCancel(rig.ctx) + r := &switchReplica{rig: rig, id: id, cancel: cancel, chatter: newTally(), + builds: map[cluster.Carrier]int{}, failBuild: map[cluster.Carrier]error{}, holdBuild: map[cluster.Carrier]chan struct{}{}} + DeferCleanup(cancel) + + cred := cluster.NewPeerCredential() + r.tunnels = tunnel.NewRegistry(rig.clusterR, id) + r.sessions = tunnel.NewPeerSessions(tunnel.NewRelay(r.tunnels).Stream) + DeferCleanup(r.sessions.Close) + r.pool = tunnel.NewPeerPool(id, cred, rig.clusterR) + DeferCleanup(r.pool.Close) + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(rig.nodeReg, r.tunnels)) + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(rig.clusterR, r.sessions.Accept)) + srv := httptest.NewServer(e) + DeferCleanup(srv.Close) + r.url = srv.URL + + r.member = cluster.NewMembership(rig.clusterR, id, "switch-spec") + r.member.SetPeer(strings.TrimPrefix(srv.URL, "http://"), cred) + r.member.SetReclaimer(r.tunnels) + Expect(r.member.Start(ctx)).To(Succeed()) + + // The replica starts on the carrier the row names. + row, err := rig.store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + first, err := r.build(ctx, row, row.Active) + Expect(err).ToNot(HaveOccurred()) + r.cur.Store(first) + r.bus = carrier.NewBroadcaster(&r.cur) + + workers := nodes.NewSwitchWorkers(rig.nodeReg, rig.clusterR, cluster.DefaultReconnectGrace, time.Hour) + r.window = carrier.NewWindow(func(ctx context.Context, nodeID string) (carrier.Attachment, error) { + cs, err := workers.AttachedCarriers(ctx, nodeID) + if err != nil { + return carrier.Attachment{}, err + } + var a carrier.Attachment + for _, c := range cs { + a.NATS = a.NATS || c == cluster.CarrierNATS + a.Tunnel = a.Tunnel || c == cluster.CarrierTunnel + } + return a, nil + }, workers.AgentsAttached) + r.commands = carrier.NewCommands(&r.cur) + r.commands.UseWindow(r.window) + r.queue = carrier.NewWorkQueue(&r.cur) + + var err2 error + r.sw, err2 = cluster.NewSwitch(cluster.SwitchOptions{ + Store: rig.store, Registry: rig.clusterR, Workers: workers, Work: rigWork{db: rig.db}, + Timings: func() cluster.Timings { return *rig.timings.Load() }, + AvailabilityMaxAge: time.Minute, + OnChange: func(row cluster.CarrierRow) { + go func() { _ = r.bus.Publish(messaging.SubjectCarrierChanged, row) }() + if hook := rig.onMove.Load(); hook != nil { + (*hook)(r, row) + } + }, + }) + Expect(err2).ToNot(HaveOccurred()) + + r.swapper, err2 = carrier.NewSwapper(carrier.SwapperOptions{ + Cur: &r.cur, Bus: r.bus, Window: r.window, Rows: rig.store, Ready: r.member, Build: r.build, + Interval: 200 * time.Millisecond, Settle: 300 * time.Millisecond, + }) + Expect(err2).ToNot(HaveOccurred()) + DeferCleanup(r.swapper.Close) + + // The chatter every replica listens to, and the dispatcher that records the + // result of a job, both through the holders. + _, err = r.bus.Subscribe(chatterSubject, r.chatter.handler) + Expect(err).ToNot(HaveOccurred()) + _, err = r.bus.Subscribe(messaging.SubjectCarrierChanged, func([]byte) { r.swapper.Wake() }) + Expect(err).ToNot(HaveOccurred()) + r.dispatch = jobs.NewDispatcher(rig.jobStore, r.queue, r.bus, rig.db, id) + Expect(r.dispatch.Start(ctx)).To(Succeed()) + DeferCleanup(r.dispatch.Stop) + + go r.swapper.Run(ctx) + go advisorylock.RunLeaderLoop(ctx, rig.db, advisorylock.KeyCarrierSwitch, 200*time.Millisecond, func() { _ = r.sw.Drive(ctx) }) + rig.replicas = append(rig.replicas, r) + return r +} + +// availability reports that every live replica can build both carriers. +func (rig *switchRig) availability() { + GinkgoHelper() + live, err := rig.clusterR.ListLive(rig.ctx, cluster.InstanceLiveness) + Expect(err).ToNot(HaveOccurred()) + for _, in := range live { + Expect(rig.clusterR.ReportAvailability(rig.ctx, in.ID, map[cluster.Carrier]string{cluster.CarrierNATS: "", cluster.CarrierTunnel: ""})).To(Succeed()) + } +} + +func (rig *switchRig) row() cluster.CarrierRow { + GinkgoHelper() + r, err := rig.store.Get(rig.ctx) + Expect(err).ToNot(HaveOccurred()) + return r +} + +// report says what the worker told the frontend about its carriers. +func (rig *switchRig) report(id string, attached []cluster.Carrier, follow []cluster.Carrier, followErr string) { + GinkgoHelper() + Expect(rig.nodeReg.SetCarrierReport(rig.ctx, id, nodes.CarrierReport{Attached: attached, Follow: follow, FollowError: followErr})).To(Succeed()) +} + +func (rig *switchRig) newJob() string { + GinkgoHelper() + task := &jobs.TaskRecord{UserID: "u1", Name: "t", Model: "m", Enabled: true} + Expect(rig.jobStore.CreateTask(task)).To(Succeed()) + job := &jobs.JobRecord{TaskID: task.ID, UserID: "u1", Status: "running", TriggeredBy: "manual"} + Expect(rig.jobStore.CreateJob(job)).To(Succeed()) + return job.ID +} + +func (rig *switchRig) enqueue(via *switchReplica, jobID string) { + GinkgoHelper() + job, err := rig.jobStore.GetJob(jobID) + Expect(err).ToNot(HaveOccurred()) + Expect(via.queue.Enqueue(rig.ctx, messaging.WorkMCPCI, jobs.JobEvent{JobID: jobID, TaskID: job.TaskID, UserID: "u1"})).To(Succeed()) +} + +func (rig *switchRig) status(jobID string) string { + job, err := rig.jobStore.GetJob(jobID) + if err != nil { + return err.Error() + } + return job.Status +} + +func (rig *switchRig) completed(ids ...string) func() bool { + return func() bool { + for _, id := range ids { + if rig.status(id) != "completed" { + return false + } + } + return true + } +} + +// request asks for a change as the admin API does: replicas have looked, then the +// preflight runs, then the row moves to prepare. +func (rig *switchRig) request(target cluster.Carrier, force bool) (cluster.CarrierRow, cluster.Report, error) { + rig.availability() + return rig.replicas[0].sw.Request(rig.ctx, cluster.Request{Target: target, By: "admin", Force: force}) +} + +// listeners counts the LISTEN sessions of the pgbus carrier on the database. +func (rig *switchRig) listeners() int64 { + GinkgoHelper() + var n int64 + Expect(rig.db.Raw("SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'").Scan(&n).Error).To(Succeed()) + return n +} + +func (rig *switchRig) eventually(f func() bool, what string) { + GinkgoHelper() + Eventually(f, "60s", "50ms").Should(BeTrue(), what) +} + +var _ = Describe("A change of carrier, end to end", Ordered, func() { + var rig *switchRig + + BeforeEach(func() { + url, err := startSwitchNATS() + if err != nil { + if os.Getenv("CI") != "" && runtime.GOOS != "darwin" { + Fail("testcontainers requires Docker and CI is set: " + err.Error()) + } + Skip("testcontainers requires Docker: " + err.Error()) + } + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + rig = &switchRig{ctx: ctx, natsURL: url, runs: newRunLog(), + installGate: make(chan struct{}), installDone: make(chan error, 1)} + rig.timings.Store(&cluster.Timings{PrepareTimeout: 20 * time.Second, TransitionWindow: 20 * time.Second, MaxDrain: 4 * time.Second}) + rig.db, rig.dsn = testutil.SetupTestDBWithDSN() + Expect(cluster.Migrate(ctx, rig.db)).To(Succeed()) + Expect(jobs.MigrateClaims(ctx, rig.db)).To(Succeed()) + rig.store, err = cluster.NewCarrierStore(rig.db) + Expect(err).ToNot(HaveOccurred()) + rig.clusterR = cluster.NewRegistry(rig.db) + rig.nodeReg, err = nodes.NewNodeRegistry(rig.db) + Expect(err).ToNot(HaveOccurred()) + rig.jobStore, err = jobs.NewJobStore(rig.db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = rig.store.Seed(ctx, cluster.CarrierNATS, "seed") + Expect(err).ToNot(HaveOccurred()) + + // The worker: a gRPC backend, an HTTP control plane for the tunnel, and + // responders on NATS for the same verbs. + rig.backend = &fakeBackend{loading: make(chan string, 1)} + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + gs := ggrpc.NewServer() + pb.RegisterBackendServer(gs, rig.backend) + go func() { _ = gs.Serve(lis) }() + DeferCleanup(gs.Stop) + rig.grpcAddr = lis.Addr().String() + rig.control = &controlPlane{address: rig.grpcAddr} + dir := GinkgoT().TempDir() + hl, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + rig.httpAddr = hl.Addr().String() + Expect(hl.Close()).To(Succeed()) + httpSrv, err := nodes.StartFileTransferServerWithControl(rig.httpAddr, dir, dir, dir, registrationToken, 1<<30, nil, nil, rig.control.handler()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { nodes.ShutdownFileTransferServer(httpSrv) }) + + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken), Address: rig.grpcAddr, HTTPAddress: rig.httpAddr} + Expect(rig.nodeReg.Register(ctx, node, true)).To(Succeed()) + rig.nodeID = node.ID + Expect(rig.nodeReg.SetTunnelTokenHash(ctx, rig.nodeID, hashOf(workerToken))).To(Succeed()) + rig.report(rig.nodeID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + + rig.workerNATS, err = messaging.New(url) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(rig.workerNATS.Close) + reply := func(v any) []byte { b, _ := json.Marshal(v); return b } + _, err = rig.workerNATS.SubscribeReply(messaging.SubjectNodeBackendInstall(rig.nodeID), func(data []byte, respond func([]byte)) { + var req workerctl.BackendInstallRequest + _ = json.Unmarshal(data, &req) + rig.natsInstalls.Add(1) + if req.OperationID == "load-slow" { + <-rig.installGate + } + respond(reply(workerctl.BackendInstallReply{Success: true, Address: rig.grpcAddr, ProcessInstance: "i1", ReportsOperations: true})) + }) + Expect(err).ToNot(HaveOccurred()) + _, err = rig.workerNATS.SubscribeReply(messaging.SubjectNodeModelOp(rig.nodeID), func(data []byte, respond func([]byte)) { + var req workerctl.OperationRequest + _ = json.Unmarshal(data, &req) + rig.natsRenewals.Add(1) + respond(reply(workerctl.OperationReply{Renewed: req.Renew, Completed: req.Complete})) + }) + Expect(err).ToNot(HaveOccurred()) + + // The agent worker: a consumer on NATS. The claim queue's consumer is the + // dispatch loop of a replica, which hands work to an agent worker that + // holds a tunnel; see followAgent. + agent := &nodes.BackendNode{Name: "agent-1", NodeType: nodes.NodeTypeAgent, TokenHash: hashOf(registrationToken)} + Expect(rig.nodeReg.Register(ctx, agent, true)).To(Succeed()) + rig.agentID = agent.ID + Expect(rig.nodeReg.SetTunnelTokenHash(ctx, rig.agentID, hashOf(agentToken))).To(Succeed()) + rig.report(rig.agentID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + agentNATS, err := messaging.New(url) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(agentNATS.Close) + _, err = messaging.NewNATSWorkConsumer(agentNATS).Consume(ctx, messaging.WorkMCPCI, 0, rig.runs.handler("nats")) + Expect(err).ToNot(HaveOccurred()) + }) + + // followWorker connects the tunnel of the worker to a replica and tells the + // frontend that it is attached to both carriers, as a worker that follows a + // change does. + followWorker := func(to *switchReplica) { + GinkgoHelper() + t, err := worker.StartTunnel(rig.ctx, worker.TunnelConfig{ + FrontendURL: to.url, NodeID: rig.nodeID, Token: func() string { return workerToken }, + Services: map[string]worker.LocalService{ + tunnel.StreamTagGRPC: func(ctx context.Context, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", rig.grpcAddr) + }, + tunnel.StreamTagHTTP: func(ctx context.Context, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", rig.httpAddr) + }, + }, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(t.Close) + Eventually(func() bool { return to.tunnels.Holds(rig.nodeID) }, "20s", "50ms").Should(BeTrue()) + rig.report(rig.nodeID, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + } + followAgent := func(to *switchReplica) *agentworker.Runtime { + GinkgoHelper() + work := agentworker.NewWork() + _, err := work.Consume(rig.ctx, messaging.WorkMCPCI, 0, rig.runs.handler("tunnel")) + Expect(err).ToNot(HaveOccurred()) + rt, err := agentworker.Start(rig.ctx, agentworker.Options{ + FrontendURL: to.url, NodeID: rig.agentID, TunnelToken: func() string { return agentToken }, + ControlToken: registrationToken, Handler: agentworker.Handler(agentworker.Config{}, work), + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(rt.Close) + Eventually(func() bool { return to.tunnels.Holds(rig.agentID) }, "20s", "50ms").Should(BeTrue()) + rig.report(rig.agentID, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + return rt + } + + It("moves a cluster with work in flight from NATS to the tunnel and back, and loses and repeats nothing", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + Expect(a.active()).To(Equal(cluster.CarrierNATS)) + + // Chatter: both replicas publish a numbered message every few milliseconds + // all through the changes, and every replica listens. + var sent sync.Map + stop := make(chan struct{}) + var chatter sync.WaitGroup + for _, r := range []*switchReplica{a, b} { + chatter.Go(func() { + for i := 0; ; i++ { + select { + case <-stop: + return + case <-time.After(5 * time.Millisecond): + } + msg := fmt.Sprintf("%s-%d", r.id, i) + if err := r.bus.Publish(chatterSubject, msg); err == nil { + sent.Store(msg, struct{}{}) + } + } + }) + } + + // Renewals of a load, once every 100 ms through the holder, with the longest + // gap between two answers. + var ( + renewMu sync.Mutex + lastOK = time.Now() + maxGap time.Duration + renewErrs []error + ) + renewStop := make(chan struct{}) + var renewing sync.WaitGroup + renewing.Go(func() { + for { + select { + case <-renewStop: + return + case <-time.After(100 * time.Millisecond): + } + _, err := a.commands.OperationControl(rig.nodeID, workerctl.OperationRequest{Renew: []string{"load-1"}}) + renewMu.Lock() + if err != nil { + renewErrs = append(renewErrs, err) + } else { + maxGap = max(maxGap, time.Since(lastOK)) + lastOK = time.Now() + } + renewMu.Unlock() + } + }) + + // An install that is under way on NATS, and a run that is held on NATS. + go func() { + _, err := a.commands.InstallBackendOp(rig.nodeID, "llama-cpp", "m", "", 0, "op-slow", "load-slow", time.Minute, nil) + rig.installDone <- err + }() + Eventually(rig.natsInstalls.Load, "10s").Should(BeEquivalentTo(1)) + heldJob := rig.newJob() + release := rig.runs.hold(heldJob) + rig.enqueue(a, heldJob) + Eventually(func() int { return len(rig.runs.ranOn(heldJob)) }, "20s", "20ms").Should(Equal(1)) + + var onNATS []string + for range 8 { + id := rig.newJob() + onNATS = append(onNATS, id) + rig.enqueue(a, id) + } + rig.eventually(rig.completed(onNATS...), "the jobs queued on NATS ran before the change") + + // The tunnel side of the worker and of the agent worker is up, and the + // preflight, which is a dry run, passes. + followAgent(a) + followWorker(b) + rig.availability() + report, err := a.sw.Preflight(rig.ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(report.OK).To(BeTrue(), "%+v", report.Blockers) + Expect(report.Replicas).To(HaveLen(2)) + Expect(report.InFlight.Jobs).To(BeNumerically(">=", 1)) + + By("changing to the tunnel with the install, the run and the renewals in flight") + began := time.Now() + _, _, err = rig.request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel + }, "the change commits and settles") + GinkgoWriter.Printf("change to the tunnel took %s\n", time.Since(began)) + rig.eventually(func() bool { return a.active() == cluster.CarrierTunnel && b.active() == cluster.CarrierTunnel }, "both replicas use the tunnel") + Expect(rig.row().Draining).To(Equal(cluster.CarrierNATS)) + + // New work goes to the claim queue, and the dispatch loop of a replica + // drives it on the agent worker that holds a tunnel. + var onTunnel []string + for range 8 { + id := rig.newJob() + onTunnel = append(onTunnel, id) + rig.enqueue(b, id) + } + rig.eventually(rig.completed(onTunnel...), "the jobs queued after the change ran on the tunnel") + + By("letting the install and the run finish where they began") + close(rig.installGate) + var installErr error + Eventually(rig.installDone, "20s").Should(Receive(&installErr)) + Expect(installErr).ToNot(HaveOccurred(), "an install that began on NATS ends on NATS") + Expect(rig.status(heldJob)).To(Equal("running")) + close(release) + rig.eventually(rig.completed(heldJob), "the run that began on NATS reports its result, which the frontend still hears on NATS") + + By("waiting for the drain to end") + rig.eventually(func() bool { return rig.row().Draining == "" }, "the leader ends the drain") + rig.eventually(func() bool { + a.mu.Lock() + defer a.mu.Unlock() + for _, c := range a.natsConns { + if c.IsConnected() { + return false + } + } + return len(a.natsConns) > 0 + }, "the old carrier is closed on a replica when the drain is over") + + close(renewStop) + renewing.Wait() + close(stop) + chatter.Wait() + renewMu.Lock() + Expect(renewErrs).To(BeEmpty(), "no renewal failed across the change") + Expect(maxGap).To(BeNumerically("<", 3*time.Second), "renewals were never missed for longer than the cadence of a lease") + renewMu.Unlock() + Expect(rig.natsRenewals.Load()).To(BeNumerically(">", 0)) + Expect(rig.control.seenRenewals()).To(BeNumerically(">", 0), "renewals went to the worker over the tunnel once the tunnel was active") + + // Every job ran once, on the carrier it was queued on. + for _, id := range onNATS { + Expect(rig.runs.ranOn(id)).To(Equal([]string{"nats"}), id) + } + Expect(rig.runs.ranOn(heldJob)).To(Equal([]string{"nats"})) + for _, id := range onTunnel { + Expect(rig.runs.ranOn(id)).To(Equal([]string{"tunnel"}), id) + } + // Every message was heard once by every replica. + time.Sleep(500 * time.Millisecond) + heardA, heardB := a.chatter.copy(), b.chatter.copy() + type span struct{ n, min, max int } + missing, repeated := map[string]*span{}, map[string]*span{} + total := 0 + note := func(into map[string]*span, key string, idx int) { + sp := into[key] + if sp == nil { + sp = &span{min: idx, max: idx} + into[key] = sp + } + sp.n++ + sp.min, sp.max = min(sp.min, idx), max(sp.max, idx) + } + sent.Range(func(k, _ any) bool { + total++ + pub, idxText, _ := strings.Cut(strings.TrimPrefix(k.(string), "replica-"), "-") + var idx int + _, _ = fmt.Sscanf(idxText, "%d", &idx) + for name, h := range map[string]map[string]int{"a": heardA, "b": heardB} { + switch n := h[k.(string)]; { + case n == 0: + note(missing, name+" missed "+pub, idx) + case n > 1: + note(repeated, name+" heard twice "+pub, idx) + } + } + return true + }) + var wrong []string + for k, sp := range missing { + wrong = append(wrong, fmt.Sprintf("%s: %d messages, index %d to %d", k, sp.n, sp.min, sp.max)) + } + for k, sp := range repeated { + wrong = append(wrong, fmt.Sprintf("%s: %d messages, index %d to %d", k, sp.n, sp.min, sp.max)) + } + slices.Sort(wrong) + Expect(total).To(BeNumerically(">", 100)) + Expect(wrong).To(BeEmpty(), "of %d messages", total) + }) + + // toTunnel changes a cluster that runs on NATS to the tunnel and waits until it + // has settled. + toTunnel := func() { + GinkgoHelper() + _, _, err := rig.request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel + }, "the change to the tunnel settles") + } + drained := func() { + GinkgoHelper() + rig.eventually(func() bool { return rig.row().Draining == "" }, "the drain ends") + } + + It("hands over the work that is still queued on the tunnel when the cluster goes back to NATS, and runs each unit once", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + agent := followAgent(a) + followWorker(a) + toTunnel() + drained() + + // No agent worker holds a tunnel any more, so the units wait in the table. + Expect(agent.Close()).To(Succeed()) + rig.report(rig.agentID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + var queued []string + for range 6 { + id := rig.newJob() + queued = append(queued, id) + rig.enqueue(a, id) + } + Consistently(rig.runs.total, "1s", "100ms").Should(BeZero()) + + _, _, err := rig.request(cluster.CarrierNATS, false) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == cluster.CarrierNATS + }, "the change back to NATS settles") + drained() + + rig.eventually(rig.completed(queued...), "the units that were queued on the tunnel run on NATS") + for _, id := range queued { + Expect(rig.runs.ranOn(id)).To(Equal([]string{"nats"}), id) + } + var migrated int64 + Expect(rig.db.Model(&jobs.WorkClaim{}).Where("state = ?", jobs.ClaimMigrated).Count(&migrated).Error).To(Succeed()) + Expect(migrated).To(BeEquivalentTo(len(queued))) + Eventually(rig.listeners, "20s", "100ms").Should(BeZero(), "the LISTEN connection is closed when the tunnel is released") + Expect(a.active()).To(Equal(cluster.CarrierNATS)) + Expect(b.active()).To(Equal(cluster.CarrierNATS)) + }) + + It("resolves a change when a replica dies in prepare, and the replica that starts again follows the row", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + c := rig.startReplica("replica-c") + // C never becomes ready: it is stuck building the target when it dies. + hold := make(chan struct{}) + c.mu.Lock() + c.holdBuild[cluster.CarrierTunnel] = hold + c.mu.Unlock() + + _, _, err := rig.request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { + in, err := rig.clusterR.Get(rig.ctx, "replica-a") + return err == nil && in.ReadyEpoch == rig.row().Epoch + }, "a replica is ready") + Expect(rig.row().State).To(Equal(cluster.StatePrepare), "C holds the change in prepare") + + c.kill() + close(hold) + rig.eventually(func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel + }, "the change commits without the replica that died") + Expect(a.active()).To(Equal(cluster.CarrierTunnel)) + Expect(b.active()).To(Equal(cluster.CarrierTunnel)) + + // The replica comes back. It reads the row, and never its own history. + again := rig.startReplica("replica-c") + Expect(again.active()).To(Equal(cluster.CarrierTunnel)) + Expect(again.buildCount(cluster.CarrierNATS)).To(BeZero(), "it never builds the carrier it left") + Expect(a.bus.Publish(chatterSubject, "after-restart")).To(Succeed()) + Eventually(func() int { return again.chatter.get("after-restart") }, "10s").Should(Equal(1), "it listens where the cluster publishes") + }) + + It("resolves a change when the replica that leads dies between the commit and the settle", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + var once sync.Once + hook := func(r *switchReplica, row cluster.CarrierRow) { + if row.State == cluster.StateCommit { + once.Do(func() { go r.kill() }) + } + } + rig.onMove.Store(&hook) + + _, _, err := rig.request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel + }, "the other replica settles the change from the row") + var survivors int + for _, r := range []*switchReplica{a, b} { + if r.active() == cluster.CarrierTunnel { + survivors++ + } + } + Expect(survivors).To(BeNumerically(">=", 1)) + }) + + It("aborts a change that is stuck in prepare, drops what was built for it, and leaves the cluster on its carrier", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + hold := make(chan struct{}) + b.mu.Lock() + b.holdBuild[cluster.CarrierTunnel] = hold + b.mu.Unlock() + + _, _, err := rig.request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { return a.buildCount(cluster.CarrierTunnel) == 1 }, "A built the target") + Eventually(rig.listeners, "10s", "50ms").Should(Equal(int64(1)), "A listens on the database for the target") + + row, err := a.sw.Abort(rig.ctx, "admin") + Expect(err).ToNot(HaveOccurred()) + Expect(row.State).To(Equal(cluster.StateStable)) + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + close(hold) + + Eventually(rig.listeners, "20s", "100ms").Should(BeZero(), "what was built for the change is dropped on both replicas") + Expect(a.active()).To(Equal(cluster.CarrierNATS)) + Expect(b.active()).To(Equal(cluster.CarrierNATS)) + + // The cluster works as it did. + id := rig.newJob() + rig.enqueue(a, id) + rig.eventually(rig.completed(id), "work still runs on NATS") + Expect(rig.runs.ranOn(id)).To(Equal([]string{"nats"})) + heardBefore := b.chatter.get("after-abort") + Expect(a.bus.Publish(chatterSubject, "after-abort")).To(Succeed()) + Eventually(func() int { return b.chatter.get("after-abort") }, "10s").Should(Equal(heardBefore + 1)) + }) + + It("blocks a change on a worker that cannot follow, and a forced change leaves it unroutable and unreaped", func() { + a, _ := rig.startReplica("replica-a"), rig.startReplica("replica-b") + rig.report(rig.nodeID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS}, "it holds no tunnel credential") + Expect(rig.nodeReg.SetNodeModel(rig.ctx, rig.nodeID, "keep-me", 0, "loaded", rig.grpcAddr, 0)).To(Succeed()) + + _, _, err := rig.request(cluster.CarrierTunnel, false) + var blocked *cluster.BlockedError + Expect(errors.As(err, &blocked)).To(BeTrue(), "%v", err) + Expect(blocked.Report.Blockers).To(ContainElement(SatisfyAll( + HaveField("Kind", cluster.BlockerWorker), + HaveField("ID", rig.nodeID), + HaveField("Reason", ContainSubstring("it holds no tunnel credential")), + ))) + Expect(rig.row().State).To(Equal(cluster.StateStable)) + Expect(rig.row().Epoch).To(Equal(int64(1))) + + _, _, err = rig.request(cluster.CarrierTunnel, true) + Expect(err).ToNot(HaveOccurred()) + rig.eventually(func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel + }, "the forced change settles") + Expect(rig.row().Force).To(BeTrue()) + + // While NATS drains the worker is still reached on it. + _, err = a.commands.OperationControl(rig.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(err).ToNot(HaveOccurred(), "the window routes to the carrier the worker is attached to") + drained() + Eventually(func() bool { + _, err := a.commands.OperationControl(rig.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + return errors.Is(err, nodes.ErrNoRoute) + }, "20s", "100ms").Should(BeTrue(), "after the drain the worker has no route, and that is all it has") + + models, err := rig.nodeReg.GetNodeModels(rig.ctx, rig.nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(HaveLen(1), "nothing was reaped") + node, err := rig.nodeReg.Get(rig.ctx, rig.nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Status).To(Equal(nodes.StatusHealthy)) + }) +}) diff --git a/core/services/carrier/tunnel.go b/core/services/carrier/tunnel.go index ea3b686a076d..e22d21f97c3a 100644 --- a/core/services/carrier/tunnel.go +++ b/core/services/carrier/tunnel.go @@ -51,6 +51,11 @@ type TunnelOptions struct { // FileManager is the shared store, and it is required with S3Staging. S3Staging bool FileManager *storage.FileManager + + // Claims makes the set claim queued work while it is in use or draining, and + // hand over what is still queued when it is released. It may be nil for a set + // that is built to be checked and thrown away. + Claims *ClaimWork } // NewTunnelSet builds the set of seam implementations that reach workers through @@ -127,6 +132,10 @@ func NewTunnelSet(o TunnelOptions) (*Set, error) { } }, } + if o.Claims != nil { + set.Start = o.Claims.start(func() *Set { return set }, o.AgentSelector, o.Token) + set.Handoff = o.Claims.handoff() + } if err := set.Validate(); err != nil { return nil, err } From 514963fdc833bfa7d390c1b95bb5bd1b6a847469 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 11:59:12 +0000 Subject: [PATCH 114/162] docs(agents): describe the change of carrier in the guide to the seams The guide says how the cluster row, the leader, the swapper and the window work together, what the preflight lists, how a run that the end of a drain cuts off is settled, why the NATS set guards loopback addresses of workers that hold a tunnel, and where the settings of the cluster live. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 104 +++++++++++++++++++++++++++++++++-- 1 file changed, 99 insertions(+), 5 deletions(-) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 4e885e110fc0..66cc296d22fc 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -37,7 +37,7 @@ call loads an `atomic.Pointer[carrier.Set]` once and calls the same method on that set. A `carrier.Set` is the implementation of every seam for one carrier. Build a set completely and call `Validate` before you store it. The holders take no lock and allocate nothing, and a call that is already running is never moved. -`carrier.NewNATSSet` is the one place that builds the NATS set. +`carrier.NewNATSSet` and `carrier.NewTunnelSet` are the places that build a set. `NodeControl` is the whole control surface a carrier gives the frontend: `NodeCommandSender`, the process lister, the three unload interfaces the model @@ -63,13 +63,107 @@ parts of that carrier before `Validate` accepts it. The active carrier is one row of table `cluster_carrier` (`cluster.CarrierStore`). Each transition is a compare-and-set on its epoch. At startup a replica reads the row, or inserts it if it is absent, and builds the -set the row names. A replica whose row names a carrier it does not implement -fails to start. It never falls back to another carrier. - -A NATS URL that points at a server which is not up does not stop the start: +set the row names (`carrierRuntime.build` in `core/application`: `NewNATSSet` or +`NewTunnelSet`). A replica whose row names a carrier it does not implement fails +to start. It never falls back to another carrier, and its own flags never +override the row. + +The row is seeded by the first replica that starts: NATS when it has a NATS URL +(the flag `--nats-url`, or a URL stored in the cluster settings), the tunnel +otherwise. `--nats-url` means what it always meant for a deployment on NATS, and +it is no longer required: a deployment with only PostgreSQL runs on the tunnel. +The flag is copied into the cluster setting `nats.url` when none is stored, and a +stored URL wins afterwards, so that every replica uses the same one. A NATS URL +that points at a server which is not up does not stop the start: `messaging.New` retries on a failed connect, so a frontend can start before its broker. A URL it cannot parse does. +## Changing the carrier + +An admin changes the carrier of a running cluster with `POST /api/cluster/carrier` +(`{"target": "nats"|"tunnel", "dry_run": bool, "force": bool}`, or `{"abort": true}`). +`GET /api/cluster/carrier` reports the state. Both routes, and +`GET`/`PUT /api/cluster/settings`, answer to an admin only; the prefix +`/api/cluster/` is not public, and only the exact connect and peer paths skip the +global authentication. Saving a NATS URL stores it and checks that the serving +replica reaches it. It does not change the carrier. + +The row has a state, and every move is a compare-and-set on the epoch: + +| State | Meaning | Who moves it on | +|---|---|---| +| `stable` | One carrier is active. A previous carrier may still be attached (`draining`, `draining_until`). | the admin starts a change | +| `prepare` | Every replica builds the target and listens on it. Publishing stays on the old carrier. | the leader, when every live replica reported ready for this epoch | +| `commit` | The target is active. Every replica publishes on it and confirms with the epoch of the commit. | the leader, when every live replica confirmed | +| `stable` | The old carrier drains for `max_drain`. | the leader, at `draining_until` | + +`cluster.Switch` holds the protocol and keeps no state outside the row and the +`instances` table, so a leader that dies between two moves is replaced by one that +reads the same row. The leader is whichever replica holds the advisory lock +`KeyCarrierSwitch` at a tick (`advisorylock.RunLeaderLoop`). Every timeout is +measured on the clock of the database (`CarrierStore.DBNow` and the change stamp), +never on the clock of a replica. A replica that cannot build the target reports +the reason, and the leader aborts naming it. A replica that is late for the prepare +timeout aborts the change, or is left behind by a forced one. `abort` is accepted +only in `prepare`: after the commit the target is the active carrier, and going +back is a change like any other. + +The preflight (also the dry run) lists: the live replicas with their versions (old +frontends write no row and cannot be detected, so the admin confirms that the list +is complete, and every frontend must be upgraded before the first change); whether +each replica could build the target (`instances.availability`, written by the +replica when a dry run asks for it); every worker that cannot follow and why +(`nodes.SwitchWorkers`: a worker that reports no capabilities predates carrier +switching, a worker that reports an error says it); and the work in flight. A +blocker can be forced. A forced change commits without a replica that is not ready, +and names it in the note of the row. + +`carrier.Swapper` is what a replica does. It polls the row every two seconds with a +jitter, and a hint on `state.carrier` makes it look at once; the poll decides. In +`prepare` it builds the target, attaches every subscription to it +(`Broadcaster.Listen`) and reports ready for the epoch. In `commit` it stores the +set, so that publishes, enqueues and new calls use it, starts the work of the set +(`Set.Start`: the claim loop of the tunnel) and confirms. In `stable` with a drain it +does nothing: the old set stays attached. When the leader ends the drain it stops +listening on the old set, cancels its work with `messaging.ErrCarrierReleased`, +runs `Set.Handoff` (the tunnel moves its pending claims to the queue of the new +carrier: update, then publish, so a job is lost and not run twice if the replica +dies in between; it runs again after `Settle` for the producers that were slow to +flip), closes the set, and runs the reconnect hooks once. A replica that polls +late, restarts or joins during a change acts on the row alone. + +`carrier.Window` is the only per-worker routing, and it exists between the commit +and the end of the drain. A control verb, a file transfer, a dial or a client for +one worker goes to the active set when the worker is attached to it, else to the +previous set when it is attached to that, else it fails with `nodes.ErrNoRoute`. +The choice is read from the state of the worker (`BackendNode.Attached`, or a +tunnel held when it reports nothing) and never from the error of an earlier +attempt. A call that is running is never moved, so an install, a stream or a load +renewal ends on the carrier it began on while the old carrier is attached. + +A run that the end of the drain cuts off is cancelled with +`messaging.ErrCarrierReleased`. The claim driver completes its claim and leaves the +job to the reaper (`ReapStuckJobs`), because the run did start and the carrier that +took over must not start it again. The model loads, installs and agent runs that +were started before the change finish where they began; whatever is still running +at `max_drain` is failed by the reaper, and a load whose lease is not renewed is +killed by the watchdog of the worker, as it is for a dead owner (`#12524`). + +A worker that holds a tunnel has no address and binds its backends to loopback. The +NATS set therefore builds its clients with `nodes.NewGuardedDirectClientFactory`: +for a worker that registered no address, a loopback address is not dialled, and the +client reports a failure of the transport, which no code that decides whether a +backend is dead reads as an answer. A forced change to NATS cannot make the health +monitor or the reconciler reap a model that runs. + +The cluster settings (`cluster.SettingsStore`, table `cluster_settings`) hold what +every replica must read the same way: `nats.url`, `nats.worker_url` and the waits +of a change (`switch.prepare_timeout`, `switch.transition_window`, +`switch.max_drain`; defaults 1m, 2m and 15m; the flags +`--carrier-prepare-timeout`, `--carrier-transition-window` and `--carrier-max-drain` +are the fallback). The runtime settings of the application are a file of one +process and cannot carry them. Credentials are never stored there. + ## The worker tunnel A worker that has no address that others can reach holds one outbound websocket From 5fb5527aa8a339b82a16aa186804236d561c4aac Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 12:00:59 +0000 Subject: [PATCH 115/162] test(carrier): fix the findings of the linter in the specs of the swap Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/swap_test.go | 13 ++++--------- core/services/cluster/switch.go | 7 ------- 2 files changed, 4 insertions(+), 16 deletions(-) diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go index 80910f3a3d4c..2fa7df780c5d 100644 --- a/core/services/carrier/swap_test.go +++ b/core/services/carrier/swap_test.go @@ -25,7 +25,6 @@ const swapSubject = "state.swap-spec" // each carrier, so a message that one replica publishes on a carrier reaches the // replicas that listen on it. type fakeNet struct { - mu sync.Mutex bus map[cluster.Carrier]*testutil.FakeBus } @@ -122,12 +121,6 @@ func (n *replicaNode) failBuild(name cluster.Carrier, err error) { func (n *replicaNode) active() cluster.Carrier { return n.cur.Load().Name } -func (n *replicaNode) handoffs() []cluster.Carrier { - n.mu.Lock() - defer n.mu.Unlock() - return slices.Clone(n.handedTo) -} - var _ = Describe("The swap of a replica", func() { var ( ctx context.Context @@ -342,8 +335,10 @@ var _ = Describe("The swap of a replica", func() { var heardOld, heardNew atomic.Int32 _, err := b.bus.Subscribe(swapSubject, func([]byte) {}) Expect(err).ToNot(HaveOccurred()) - net.bus[cluster.CarrierNATS].Subscribe(swapSubject, func([]byte) { heardOld.Add(1) }) - net.bus[cluster.CarrierTunnel].Subscribe(swapSubject, func([]byte) { heardNew.Add(1) }) + _, err = net.bus[cluster.CarrierNATS].Subscribe(swapSubject, func([]byte) { heardOld.Add(1) }) + Expect(err).ToNot(HaveOccurred()) + _, err = net.bus[cluster.CarrierTunnel].Subscribe(swapSubject, func([]byte) { heardNew.Add(1) }) + Expect(err).ToNot(HaveOccurred()) request(cluster.CarrierTunnel) poll() diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index d5d81482723a..944f758bf842 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -263,13 +263,6 @@ func (s *Switch) observePhase(phase string, d time.Duration) { } } -func other(c Carrier) Carrier { - if c == CarrierNATS { - return CarrierTunnel - } - return CarrierNATS -} - // Status reports the cluster as it is, with no target. func (s *Switch) Status(ctx context.Context) (Report, error) { return s.report(ctx, "") From 21d922e834767f84b6cb2619671d7b1c40274979 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 12:02:12 +0000 Subject: [PATCH 116/162] fix(carrier): mark the jitter of the poll as not security relevant Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/swap.go | 1 + 1 file changed, 1 insertion(+) diff --git a/core/services/carrier/swap.go b/core/services/carrier/swap.go index 066d42beebd8..1afc6affbcb8 100644 --- a/core/services/carrier/swap.go +++ b/core/services/carrier/swap.go @@ -216,6 +216,7 @@ func (s *Swapper) Run(ctx context.Context) { } else if err == nil { lastErr = "" } + // #nosec G404 -- spreads the polls of the replicas apart; the value is no secret. wait := time.Duration(float64(s.o.Interval) * (0.8 + 0.4*rand.Float64())) timer := time.NewTimer(wait) select { From 1b3ba02631fa98146e10ec925173663ed1402808 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 12:35:15 +0000 Subject: [PATCH 117/162] test(carrier): check the metrics of the swap and the status of the switch The metrics of a replica show the epoch of the row, the state of a change, the carrier in use, whether the replica is ready, and how long the drain has to go. The status of the switch reports the row, the replicas, the workers and the time left of a drain without a target, and blocks nothing. A deployment with no NATS URL has no NATS bus, so the warning about a bus without credentials is no longer printed for it. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/config/distributed_config.go | 6 ++- core/services/carrier/swap_test.go | 76 ++++++++++++++++++++++++++++ core/services/cluster/switch_test.go | 27 ++++++++++ 3 files changed, 108 insertions(+), 1 deletion(-) diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 4d93e266032a..82e8e8202a3a 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -188,7 +188,11 @@ func (c DistributedConfig) Validate() error { if err := c.NatsTLSFiles().Validate(); err != nil { return err } - c.NatsAuthConfig().WarnIfInsecure(true) + // A deployment with no NATS URL runs on the tunnel and has no NATS bus to warn + // about. The warning comes back when a URL is configured. + if c.NatsURL != "" { + c.NatsAuthConfig().WarnIfInsecure(true) + } // Check for negative durations for name, d := range map[string]time.Duration{ FlagMCPToolTimeout: c.MCPToolTimeout, diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go index 2fa7df780c5d..417d9a8a84ba 100644 --- a/core/services/carrier/swap_test.go +++ b/core/services/carrier/swap_test.go @@ -15,6 +15,9 @@ import ( "github.com/mudler/LocalAI/core/services/testutil" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" "gorm.io/gorm" ) @@ -618,6 +621,69 @@ var _ = Describe("The swap of a replica", func() { }) }) + Describe("the metrics", func() { + It("report the epoch, the state, the carrier in use, the readiness and how long the drain has to go", func() { + reader := sdkmetric.NewManualReader() + meter := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)).Meter("spec") + a := join("a") + measured, err := carrier.NewSwapper(carrier.SwapperOptions{ + Cur: &a.cur, Bus: a.bus, Window: a.window, Rows: store, + Ready: readyReporter{reg: reg, id: "a"}, Build: a.build, + Interval: time.Hour, Meter: meter, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(measured.Close) + + gauge := func(name, attr string) float64 { + GinkgoHelper() + var rm metricdata.ResourceMetrics + Expect(reader.Collect(ctx, &rm)).To(Succeed()) + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != name { + continue + } + switch d := m.Data.(type) { + case metricdata.Gauge[int64]: + for _, p := range d.DataPoints { + if attr == "" || p.Attributes.Len() > 0 && attrValue(p.Attributes, attr) { + return float64(p.Value) + } + } + case metricdata.Gauge[float64]: + for _, p := range d.DataPoints { + return p.Value + } + } + } + } + Fail("no data point for " + name + " " + attr) + return 0 + } + + Expect(measured.Poll(ctx)).To(Succeed()) + Expect(gauge("localai_carrier_epoch", "")).To(BeEquivalentTo(1)) + Expect(gauge("localai_carrier_state", "stable")).To(BeEquivalentTo(1)) + Expect(gauge("localai_carrier_state", "prepare")).To(BeZero()) + Expect(gauge("localai_carrier_active", "nats")).To(BeEquivalentTo(1)) + Expect(gauge("localai_carrier_active", "tunnel")).To(BeZero()) + + request(cluster.CarrierTunnel) + Expect(measured.Poll(ctx)).To(Succeed()) + Expect(gauge("localai_carrier_state", "prepare")).To(BeEquivalentTo(1)) + Expect(gauge("localai_carrier_replica_ready", "")).To(BeEquivalentTo(1)) + drive() + Expect(measured.Poll(ctx)).To(Succeed()) + Expect(gauge("localai_carrier_state", "commit")).To(BeEquivalentTo(1)) + Expect(gauge("localai_carrier_active", "tunnel")).To(BeEquivalentTo(1)) + ready := func() { Expect(reg.ReportReady(ctx, "a", row().Epoch, "")).To(Succeed()) } + ready() + drive() + Expect(measured.Poll(ctx)).To(Succeed()) + Expect(gauge("localai_carrier_drain_remaining_seconds", "")).To(BeNumerically(">", 0), "the drain has time to go") + }) + }) + Describe("the polling loop", func() { It("follows the row on its own, and wakes at once on a hint", func() { a := join("a") @@ -655,3 +721,13 @@ type readyReporter struct { func (r readyReporter) ReportReady(ctx context.Context, epoch int64, reason string) error { return r.reg.ReportReady(ctx, r.id, epoch, reason) } + +// attrValue reports whether an attribute set holds the value. +func attrValue(set attribute.Set, want string) bool { + for _, kv := range set.ToSlice() { + if kv.Value.AsString() == want { + return true + } + } + return false +} diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go index ae393ae62c19..6f1262855f73 100644 --- a/core/services/cluster/switch_test.go +++ b/core/services/cluster/switch_test.go @@ -309,6 +309,33 @@ var _ = Describe("The switch of the carrier", func() { }) }) + Describe("the status", func() { + It("reports the row, the replicas and the workers without a target, and the time the drain has to go", func() { + replica("a", "", "") + workers.set(cluster.WorkerInfo{ID: "w1", Name: "w1", Attached: []cluster.Carrier{cluster.CarrierNATS}, Reports: true, Follow: []cluster.Carrier{cluster.CarrierNATS}}) + status, err := sw.Status(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(status.Active).To(Equal(cluster.CarrierNATS)) + Expect(status.Target).To(BeEmpty()) + Expect(status.Blockers).To(BeEmpty(), "a status is not a dry run") + Expect(status.Replicas).To(HaveLen(1)) + Expect(status.Workers).To(HaveLen(1)) + Expect(status.DrainRemaining).To(BeZero()) + + _, _, err = request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + status, err = sw.Status(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(status.Row.Draining).To(Equal(cluster.CarrierNATS)) + Expect(status.DrainRemaining).To(BeNumerically(">", 0)) + Expect(status.DrainRemaining).To(BeNumerically("<=", timings.MaxDrain)) + }) + }) + Describe("the hint", func() { It("tells the hook about every move of the row, and about no move that failed", func() { var epochs []int64 From ffb2071d379e55ce4cf4effddabc192c0529e920 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:07:47 +0000 Subject: [PATCH 118/162] fix(carrier): forward UnloadRemoteModel to the method of the same name The window routes a stop of a model by node, and the holder must still forward each method to the method of the same name on the set. The specs of the holders and of the window said so and were out of date after the first version of the window. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/holders.go | 9 ++++++++- core/services/carrier/swap_test.go | 2 +- core/services/carrier/window_test.go | 2 +- core/services/cluster/switch_test.go | 2 +- 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/core/services/carrier/holders.go b/core/services/carrier/holders.go index 4978de3fec64..2164cb5ab19b 100644 --- a/core/services/carrier/holders.go +++ b/core/services/carrier/holders.go @@ -149,8 +149,15 @@ func (h *Commands) ListRunningModels(nodeID string) (*workerctl.ModelsRunningRep return s.Commands.ListRunningModels(nodeID) } +// UnloadRemoteModel stops a model on every node that holds it, as +// UnloadRemoteModelContext does. func (h *Commands) UnloadRemoteModel(modelName string) error { - return h.UnloadRemoteModelContext(context.Background(), modelName, false) + cur := h.cur.Load() + err := cur.Commands.UnloadRemoteModel(modelName) + if prev := h.previous(cur); prev != nil && err != nil { + return prev.Commands.UnloadRemoteModel(modelName) + } + return err } // previous returns the set that is draining, or nil outside a change of carrier. diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go index 417d9a8a84ba..b88756d4c21f 100644 --- a/core/services/carrier/swap_test.go +++ b/core/services/carrier/swap_test.go @@ -587,7 +587,7 @@ var _ = Describe("The swap of a replica", func() { a := join("a") settle(cluster.CarrierTunnel) Expect(a.cmds.UnloadRemoteModel("m")).To(Succeed()) - Expect(a.fake(cluster.CarrierTunnel).commands.seen()).To(ContainElement("UnloadRemoteModelContext")) + Expect(a.fake(cluster.CarrierTunnel).commands.seen()).To(ContainElement("UnloadRemoteModel")) Expect(a.fake(cluster.CarrierNATS).commands.seen()).To(BeEmpty(), "the active carrier did it all") }) }) diff --git a/core/services/carrier/window_test.go b/core/services/carrier/window_test.go index 898618ad6da4..b627e8fa4f4b 100644 --- a/core/services/carrier/window_test.go +++ b/core/services/carrier/window_test.go @@ -128,7 +128,7 @@ var _ = Describe("The window of a change of carrier", func() { cmds := carrier.NewCommands(&cur) cmds.UseWindow(window) Expect(cmds.UnloadRemoteModel("m")).To(Succeed()) - Expect(next.commands.seen()).To(Equal([]string{"UnloadRemoteModelContext"})) + Expect(next.commands.seen()).To(Equal([]string{"UnloadRemoteModel"})) Expect(old.commands.seen()).To(BeEmpty(), "the active carrier did everything, so the previous one is not asked") }) diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go index 6f1262855f73..c13a68344e63 100644 --- a/core/services/cluster/switch_test.go +++ b/core/services/cluster/switch_test.go @@ -312,7 +312,7 @@ var _ = Describe("The switch of the carrier", func() { Describe("the status", func() { It("reports the row, the replicas and the workers without a target, and the time the drain has to go", func() { replica("a", "", "") - workers.set(cluster.WorkerInfo{ID: "w1", Name: "w1", Attached: []cluster.Carrier{cluster.CarrierNATS}, Reports: true, Follow: []cluster.Carrier{cluster.CarrierNATS}}) + workers.set(cluster.WorkerInfo{ID: "w1", Name: "w1", Attached: []cluster.Carrier{cluster.CarrierNATS}, Reports: true, Follow: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}}) status, err := sw.Status(ctx) Expect(err).ToNot(HaveOccurred()) Expect(status.Active).To(Equal(cluster.CarrierNATS)) From e631d49cdc78be79fbe6475579b092f4db733f9f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:34:54 +0000 Subject: [PATCH 119/162] fix(messaging): store the payload size of the work queue spec as 64 bits The spec stored the length of a payload in an int32, which gosec reports as an integer overflow conversion. A 64-bit counter holds any length, and the spec reads it as an int as before. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/messaging/messagingtest/workqueue.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/services/messaging/messagingtest/workqueue.go b/core/services/messaging/messagingtest/workqueue.go index 2b1d2b767c5a..2f4f23aa5132 100644 --- a/core/services/messaging/messagingtest/workqueue.go +++ b/core/services/messaging/messagingtest/workqueue.go @@ -223,9 +223,9 @@ func RunWorkQueueConformance(newRig WorkQueueFactory) { It("accepts a payload up to the bound and refuses one above it", func() { var got atomic.Int32 - var size atomic.Int32 + var size atomic.Int64 consume(rig.NewConsumer(), messaging.WorkMCPCI, 0, func(_ context.Context, p []byte, _ messaging.Publisher) error { - size.Store(int32(len(p))) + size.Store(int64(len(p))) got.Add(1) return nil }) From bbec70353e5e062abde2b0b5b0273844cbecb886 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:35:53 +0000 Subject: [PATCH 120/162] feat(messaging): accept the NATS CA as PEM in memory A worker that the frontend hands the address of NATS to has no file for the CA of that server. TLSFiles takes the CA as a PEM bundle as well as a path, validates that it holds a certificate, and trusts it, together with the file when both are set. The specs connect to a server of a private CA over TLS with the bundle, and show that the same server is refused without it. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/messaging/tls.go | 49 +++++++++++- core/services/messaging/tls_test.go | 119 ++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+), 1 deletion(-) diff --git a/core/services/messaging/tls.go b/core/services/messaging/tls.go index b594845cc197..f99ed2ea79c4 100644 --- a/core/services/messaging/tls.go +++ b/core/services/messaging/tls.go @@ -1,6 +1,9 @@ package messaging import ( + "crypto/tls" + "crypto/x509" + "errors" "fmt" "os" @@ -13,11 +16,15 @@ type TLSFiles struct { CA string // LOCALAI_NATS_TLS_CA — private CA for server verification Cert string // LOCALAI_NATS_TLS_CERT — client certificate (mTLS) Key string // LOCALAI_NATS_TLS_KEY — client private key + // CAPEM is a private CA in memory, as a PEM bundle. A worker that the + // frontend hands the CA to has no file for it. A path in CA is used as well + // when both are set. + CAPEM []byte } // Enabled reports whether any TLS file path is configured. func (f TLSFiles) Enabled() bool { - return f.CA != "" || f.Cert != "" || f.Key != "" + return f.CA != "" || f.Cert != "" || f.Key != "" || len(f.CAPEM) > 0 } // Validate checks path pairing and that files exist. @@ -28,6 +35,11 @@ func (f TLSFiles) Validate() error { if f.Key != "" && f.Cert == "" { return fmt.Errorf("LOCALAI_NATS_TLS_CERT is required when LOCALAI_NATS_TLS_KEY is set") } + if len(f.CAPEM) > 0 { + if _, err := caPool(f.CAPEM); err != nil { + return err + } + } for _, path := range []struct { name, path string }{ @@ -54,6 +66,32 @@ func (f TLSFiles) natsOptions() ([]nats.Option, error) { if f.CA != "" { opts = append(opts, nats.RootCAs(f.CA)) } + if len(f.CAPEM) > 0 { + pool, err := caPool(f.CAPEM) + if err != nil { + return nil, err + } + opts = append(opts, func(o *nats.Options) error { + previous := o.RootCAsCB + o.RootCAsCB = func() (*x509.CertPool, error) { + if previous == nil { + return pool, nil + } + // A file was given as well: trust both. + both, err := previous() + if err != nil { + return nil, err + } + both.AppendCertsFromPEM(f.CAPEM) + return both, nil + } + if o.TLSConfig == nil { + o.TLSConfig = &tls.Config{MinVersion: tls.VersionTLS12} + } + o.Secure = true + return nil + }) + } if f.Cert != "" { opts = append(opts, nats.ClientCert(f.Cert, f.Key)) } @@ -66,3 +104,12 @@ func WithTLS(files TLSFiles) Option { c.tls = files } } + +// caPool parses a PEM bundle of CA certificates. +func caPool(pem []byte) (*x509.CertPool, error) { + pool := x509.NewCertPool() + if !pool.AppendCertsFromPEM(pem) { + return nil, errors.New("the NATS CA given as PEM holds no certificate") + } + return pool, nil +} diff --git a/core/services/messaging/tls_test.go b/core/services/messaging/tls_test.go index 9fdb3e81686f..4456d293b25b 100644 --- a/core/services/messaging/tls_test.go +++ b/core/services/messaging/tls_test.go @@ -1,8 +1,20 @@ package messaging_test import ( + "bufio" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "math/big" + "net" "os" "path/filepath" + "strings" + "time" "github.com/mudler/LocalAI/core/services/messaging" @@ -23,3 +35,110 @@ var _ = Describe("TLSFiles", func() { Expect((messaging.TLSFiles{CA: ca}).Validate()).To(Succeed()) }) }) + +// fakeTLSNATS speaks just enough of the NATS protocol for a client to connect +// over TLS: it sends INFO with tls_required, upgrades the connection, and +// answers CONNECT and PING. It lets a spec check which certificates a client +// trusts without a NATS server. +type fakeTLSNATS struct { + addr string + ca []byte +} + +func startFakeTLSNATS() *fakeTLSNATS { + GinkgoHelper() + caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + Expect(err).ToNot(HaveOccurred()) + caTmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "spec ca"}, + NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), + IsCA: true, BasicConstraintsValid: true, KeyUsage: x509.KeyUsageCertSign, + } + caDER, err := x509.CreateCertificate(rand.Reader, caTmpl, caTmpl, &caKey.PublicKey, caKey) + Expect(err).ToNot(HaveOccurred()) + caCert, err := x509.ParseCertificate(caDER) + Expect(err).ToNot(HaveOccurred()) + + leafKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + Expect(err).ToNot(HaveOccurred()) + leafTmpl := &x509.Certificate{ + SerialNumber: big.NewInt(2), Subject: pkix.Name{CommonName: "127.0.0.1"}, + NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, + KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + } + leafDER, err := x509.CreateCertificate(rand.Reader, leafTmpl, caCert, &leafKey.PublicKey, caKey) + Expect(err).ToNot(HaveOccurred()) + tlsCfg := &tls.Config{Certificates: []tls.Certificate{{Certificate: [][]byte{leafDER}, PrivateKey: leafKey}}} + + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(lis.Close) + go func() { + for { + raw, err := lis.Accept() + if err != nil { + return + } + go func() { + defer func() { _ = raw.Close() }() + _, _ = raw.Write([]byte(`INFO {"server_id":"fake","version":"2.10.0","proto":1,"host":"127.0.0.1","max_payload":1048576,"tls_required":true}` + "\r\n")) + conn := tls.Server(raw, tlsCfg) + if conn.Handshake() != nil { + return + } + rd := bufio.NewReader(conn) + for { + line, err := rd.ReadString('\n') + if err != nil { + return + } + switch { + case strings.HasPrefix(line, "PING"): + _, _ = conn.Write([]byte("PONG\r\n")) + case strings.HasPrefix(line, "CONNECT"): + // The client sends PING next and waits for PONG. + } + } + }() + } + }() + return &fakeTLSNATS{ + addr: "tls://" + lis.Addr().String(), + ca: pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caDER}), + } +} + +var _ = Describe("A CA given as PEM", func() { + connects := func(url string, opts ...messaging.Option) bool { + GinkgoHelper() + c, err := messaging.New(url, opts...) + if err != nil { + return false + } + defer c.Close() + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if c.IsConnected() { + return true + } + time.Sleep(20 * time.Millisecond) + } + return false + } + + It("is trusted without a file on disk", func() { + srv := startFakeTLSNATS() + Expect(connects(srv.addr, messaging.WithTLS(messaging.TLSFiles{CAPEM: srv.ca}))).To(BeTrue()) + }) + + It("is required: a server of a private CA is refused without it", func() { + srv := startFakeTLSNATS() + Expect(connects(srv.addr)).To(BeFalse()) + }) + + It("counts as TLS configuration, and is refused when it holds no certificate", func() { + Expect((messaging.TLSFiles{CAPEM: []byte("x")}).Enabled()).To(BeTrue()) + Expect((messaging.TLSFiles{CAPEM: []byte("not a certificate")}).Validate()).To(HaveOccurred()) + }) +}) From b32b4c344e327382c39f0329265d8574029f5a9a Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:40:07 +0000 Subject: [PATCH 121/162] feat(http): hand a worker what it needs to follow a change of carrier A worker that follows a change of carrier registers again to get the credential of the carrier it attaches to. The registration takes the carrier that the worker asks for. The frontend gives the credential of that carrier when it is the active one, the target of a change under way, or the one that drains, and the credential of the active carrier for anything else, so that a node cannot get a tunnel while the cluster has none. Only that credential is minted: a worker that holds the tunnel and registers for NATS keeps its tunnel credential and its session. The answer names the carrier, the epoch, the state of a change, its target and the drain, and the carrier that the credential is for. A credential for NATS comes with the address that workers are told to use, the CA of the server as PEM, and a flag when the server asks for a client certificate, which cannot be handed over. A worker needs no NATS setting of its own for this. The answer to a heartbeat carries the same news, because the heartbeat reaches the frontend whichever carrier is active. A heartbeat can carry the carriers that the worker is attached to, the epoch it saw, the carriers it can follow and why it cannot follow another. They are stored for the frontends to route by, and a worker that sends none is left as it was. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/carrier_runtime.go | 38 +++ core/application/nats_handover_test.go | 70 +++++ core/http/app.go | 1 + core/http/endpoints/localai/nodes.go | 154 +++++++++- .../endpoints/localai/nodes_follow_test.go | 280 ++++++++++++++++++ core/http/routes/nodes.go | 2 +- core/services/nodes/follow_report.go | 28 ++ core/services/nodes/registry.go | 25 ++ 8 files changed, 585 insertions(+), 13 deletions(-) create mode 100644 core/application/nats_handover_test.go create mode 100644 core/http/endpoints/localai/nodes_follow_test.go diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go index fe8649c5cfbb..357576249302 100644 --- a/core/application/carrier_runtime.go +++ b/core/application/carrier_runtime.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "os" "sync/atomic" "time" @@ -81,6 +82,43 @@ func (rt *carrierRuntime) natsURL(ctx context.Context) (string, error) { return rt.cfg.Distributed.NatsURL, nil } +// WorkerURL is the address of the NATS server that workers are told to use: the +// stored address for workers, else the address of the frontends, else the flag of +// this replica. A worker may reach the server by another name than a frontend +// does. It is empty when the deployment has no NATS. +func (rt *carrierRuntime) WorkerURL(ctx context.Context) (string, error) { + v, ok, err := rt.settings.Get(ctx, cluster.SettingNATSWorkerURL) + if err != nil { + return "", err + } + if ok && v != "" { + return v, nil + } + return rt.natsURL(ctx) +} + +// CAPEM is the private CA of the NATS server as PEM, or empty when the server has +// none. A worker is handed it, so that it needs no file of its own to verify the +// server. +func (rt *carrierRuntime) CAPEM() (string, error) { + path := rt.cfg.Distributed.NatsTLSCA + if path == "" { + return "", nil + } + raw, err := os.ReadFile(path) + if err != nil { + return "", fmt.Errorf("reading the NATS CA to hand to workers: %w", err) + } + return string(raw), nil +} + +// ClientTLS says that this replica connects to NATS with a client certificate. A +// worker cannot be handed such a file, so it needs its own. +func (rt *carrierRuntime) ClientTLS() bool { + d := rt.cfg.Distributed + return d.NatsTLSCert != "" || d.NatsTLSKey != "" +} + // Build builds the set of a carrier for a change. It waits for the carrier to be // usable, because a replica that reports ready must be. func (rt *carrierRuntime) Build(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*carrier.Set, error) { diff --git a/core/application/nats_handover_test.go b/core/application/nats_handover_test.go new file mode 100644 index 000000000000..a4f27be56207 --- /dev/null +++ b/core/application/nats_handover_test.go @@ -0,0 +1,70 @@ +package application + +import ( + "context" + "os" + "path/filepath" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("What workers are told about NATS", func() { + var ( + ctx context.Context + settings *cluster.SettingsStore + cfg *config.ApplicationConfig + rt *carrierRuntime + ) + + BeforeEach(func() { + ctx = context.Background() + var err error + settings, err = cluster.NewSettingsStore(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + cfg = &config.ApplicationConfig{Distributed: config.DistributedConfig{NatsURL: "nats://frontend-flag:4222"}} + rt = &carrierRuntime{cfg: cfg, settings: settings} + }) + + It("is the address for workers when one is stored", func() { + Expect(settings.Set(ctx, cluster.SettingNATSURL, "nats://frontends:4222", "spec")).To(Succeed()) + Expect(settings.Set(ctx, cluster.SettingNATSWorkerURL, "nats://workers.example:4222", "spec")).To(Succeed()) + Expect(rt.WorkerURL(ctx)).To(Equal("nats://workers.example:4222")) + }) + + It("falls back to the address of the frontends, then to the flag of this replica", func() { + Expect(rt.WorkerURL(ctx)).To(Equal("nats://frontend-flag:4222")) + Expect(settings.Set(ctx, cluster.SettingNATSURL, "nats://frontends:4222", "spec")).To(Succeed()) + Expect(rt.WorkerURL(ctx)).To(Equal("nats://frontends:4222")) + }) + + It("is empty when the deployment has no NATS", func() { + cfg.Distributed.NatsURL = "" + Expect(rt.WorkerURL(ctx)).To(BeEmpty()) + }) + + It("hands over the CA of the server and says nothing when there is none", func() { + Expect(rt.CAPEM()).To(BeEmpty()) + path := filepath.Join(GinkgoT().TempDir(), "ca.pem") + Expect(os.WriteFile(path, []byte("-----BEGIN CERTIFICATE-----\nx\n-----END CERTIFICATE-----\n"), 0o600)).To(Succeed()) + cfg.Distributed.NatsTLSCA = path + Expect(rt.CAPEM()).To(ContainSubstring("BEGIN CERTIFICATE")) + }) + + It("fails to hand over a CA file it cannot read, and does not make up one", func() { + cfg.Distributed.NatsTLSCA = filepath.Join(GinkgoT().TempDir(), "missing.pem") + _, err := rt.CAPEM() + Expect(err).To(HaveOccurred()) + }) + + It("says that workers need a certificate when this replica holds one", func() { + Expect(rt.ClientTLS()).To(BeFalse()) + cfg.Distributed.NatsTLSCert = "/some/client.pem" + cfg.Distributed.NatsTLSKey = "/some/client.key" + Expect(rt.ClientTLS()).To(BeTrue()) + }) +}) diff --git a/core/http/app.go b/core/http/app.go index fe9ac2091ba4..e360a97dc1fd 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -572,6 +572,7 @@ func API(application *application.Application) (*echo.Echo, error) { tunnels = d.Tunnels registerOpts = append(registerOpts, localai.WithCarrierReader(d.Carriers)) registerOpts = append(registerOpts, localai.WithTunnelDisconnector(d)) + registerOpts = append(registerOpts, localai.WithNATSHandover(d.Runtime)) if d.Router != nil { remoteUnloader = d.Router.Unloader() } diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index ba1f09c60f3c..ce179b68ee84 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -121,6 +121,14 @@ type RegisterNodeRequest struct { // none sends false: it holds an outbound tunnel and nothing dials it. A // worker that predates the field sends nothing, and that means true. Routable *bool `json:"routable,omitempty"` + // Carrier is the carrier the worker asks the credential of. A worker that + // follows a change of carrier registers again to get the credential of the + // carrier it attaches to, which is not the active one yet while the change is + // being prepared. Only the active carrier, the target of a change under way + // and the carrier that drains can be asked for. Anything else, and a worker + // that sends nothing (it predates the field, or it boots), gets the credential + // of the active carrier. + Carrier string `json:"carrier,omitempty"` } // CarrierReader reads the row that names the carrier of the cluster. @@ -132,6 +140,30 @@ type CarrierReader interface { type registerOptions struct { carriers CarrierReader tunnels TunnelDisconnector + nats NATSHandover +} + +// NATSHandover says what a worker needs to reach NATS besides its own credential. +// The frontend hands it over in the answer to a registration, so that a worker +// that was never given the address of NATS, or that booted on the tunnel, can +// attach to NATS when the cluster changes to it. +type NATSHandover interface { + // WorkerURL is the address that workers use. It is empty when no NATS is + // configured. + WorkerURL(ctx context.Context) (string, error) + // CAPEM is the CA of the NATS server as a PEM bundle, or empty when the + // server does not use a private CA. + CAPEM() (string, error) + // ClientTLS says that the server asks for a client certificate. Files of that + // kind are not handed over, so a worker needs its own. + ClientTLS() bool +} + +// WithNATSHandover makes a registration that hands over a NATS credential also +// hand over the address of the server, its CA, and whether the worker needs a +// client certificate of its own. +func WithNATSHandover(h NATSHandover) RegisterOption { + return func(o *registerOptions) { o.nats = h } } // TunnelDisconnector ends the tunnel that this replica holds for a node. @@ -188,6 +220,42 @@ func (o registerOptions) activeCarrier(ctx context.Context) (cluster.CarrierRow, return row, true } +// credentialCarrier is the carrier whose credential an answer carries: the one +// the worker asked for when it may have it, else the active one. +func credentialCarrier(row cluster.CarrierRow, asked string) cluster.Carrier { + c := cluster.Carrier(asked) + switch { + case c == "" || c == row.Active: + return row.Active + case row.State != cluster.StateStable && c == row.Target: + return c + case c == row.Draining: + return c + } + return row.Active +} + +// carrierNews is what a worker is told about the carrier of the cluster, in the +// answer to a registration and to a heartbeat. The heartbeat is the channel that +// exists whichever carrier is active, so it is the one that carries the news. +func carrierNews(row cluster.CarrierRow) map[string]any { + view := map[string]any{ + "carrier": row.Active, + "carrier_epoch": row.Epoch, + "carrier_state": row.State, + } + if row.Target != "" { + view["carrier_target"] = row.Target + } + if row.Draining != "" { + view["carrier_draining"] = row.Draining + if row.DrainingUntil != nil { + view["carrier_draining_until"] = row.DrainingUntil.UTC().Format(time.RFC3339) + } + } + return view +} + // RegisterNodeEndpoint registers a new backend node. // expectedToken is the registration token configured on the frontend (may be empty to disable auth). // autoApprove controls whether new nodes go directly to "healthy" or require admin approval. @@ -224,7 +292,14 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au ctx := c.Request().Context() carrierRow, haveCarrier := opts.activeCarrier(ctx) - tunnelActive := haveCarrier && carrierRow.Active == cluster.CarrierTunnel + // The carrier whose credential this answer carries. It decides whether + // the address of the worker matters: a worker that asks for NATS while + // the tunnel is still active is about to be dialled. + credFor := cluster.CarrierNATS + if haveCarrier { + credFor = credentialCarrier(carrierRow, req.Carrier) + } + tunnelActive := haveCarrier && credFor == cluster.CarrierTunnel // A worker that said that its address cannot be reached holds an // outbound tunnel and nothing dials it. Its address is dropped and not // stored: a dialable-looking endpoint that nothing may dial misleads. @@ -349,7 +424,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } - attachCarrierCredential(ctx, response, registry, node, natsCfg, carrierRow, haveCarrier, opts.tunnels) + attachCarrierCredential(ctx, response, registry, node, natsCfg, carrierRow, haveCarrier, credFor, opts) if tunnelActive && expectedToken == "" && autoApprove { openRegistrationWarning.Do(func() { xlog.Warn("Worker registration is open: there is no registration token and new nodes are approved at once, " + @@ -407,19 +482,53 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr } // attachCarrierCredential names the active carrier in the answer and adds the -// credential for it: the tunnel token when the tunnel is active, the NATS user -// JWT otherwise. A deployment with no carrier row is a deployment with NATS and -// gets the answer that it always had. -func attachCarrierCredential(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, natsCfg natsauth.Config, row cluster.CarrierRow, haveCarrier bool, tunnels TunnelDisconnector) { +// credential for credFor: the tunnel token when that is the tunnel, the NATS user +// JWT and the address of the server otherwise. A deployment with no carrier row +// is a deployment with NATS and gets the answer that it always had. +// +// A credential is minted for one carrier only. A tunnel token replaces the +// previous one and ends the session that the replica holds for the node, so a +// worker that holds the tunnel and registers for NATS must not mint one. +func attachCarrierCredential(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, natsCfg natsauth.Config, row cluster.CarrierRow, haveCarrier bool, credFor cluster.Carrier, opts registerOptions) { if haveCarrier { - response["carrier"] = row.Active - response["carrier_epoch"] = row.Epoch + for k, v := range carrierNews(row) { + response[k] = v + } + response["credential_for"] = credFor } - if haveCarrier && row.Active == cluster.CarrierTunnel { - attachTunnelToken(ctx, response, registry, node, tunnels) + if haveCarrier && credFor == cluster.CarrierTunnel { + attachTunnelToken(ctx, response, registry, node, opts.tunnels) return } attachNatsJWT(response, node, natsCfg) + attachNATSHandover(ctx, response, opts.nats) +} + +// attachNATSHandover adds the address of the NATS server and what is needed to +// verify it. A failure is logged and the answer goes out without it: the worker +// then uses a URL of its own, or says that it cannot attach. +func attachNATSHandover(ctx context.Context, response map[string]any, h NATSHandover) { + if h == nil { + return + } + url, err := h.WorkerURL(ctx) + if err != nil { + xlog.Warn("Reading the NATS address for a registration failed", "error", err) + return + } + if url == "" { + return + } + response["nats_url"] = url + ca, err := h.CAPEM() + if err != nil { + xlog.Warn("Reading the NATS CA for a registration failed", "error", err) + } else if ca != "" { + response["nats_ca"] = ca + } + if h.ClientTLS() { + response["nats_client_tls"] = true + } } // attachTunnelToken mints a new tunnel credential for the node, stores only its @@ -597,7 +706,14 @@ func DeactivateNodeEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { } // HeartbeatEndpoint updates the heartbeat for a node. -func HeartbeatEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { +// +// The answer tells the worker which carrier the cluster uses and whether a change +// is under way. A worker follows the carrier from this answer, because the +// heartbeat reaches the frontend whichever carrier is active. A worker that +// reports the carriers it is attached to and the ones it can follow has them +// stored, and the frontends route by them during a change. +func HeartbeatEndpoint(registry *nodes.NodeRegistry, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) return func(c echo.Context) error { id := c.Param("id") @@ -622,7 +738,21 @@ func HeartbeatEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { xlog.Warn("Heartbeat failed for node", "id", id, "error", err) return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - return c.JSON(http.StatusOK, map[string]string{"message": "heartbeat received"}) + if report, ok := update.CarrierReport(); ok { + if err := registry.SetCarrierReport(ctx, id, report); err != nil { + xlog.Warn("Failed to store the carrier report of a node", "id", id, "error", err) + } + } + response := map[string]any{"message": "heartbeat received"} + if row, ok := opts.activeCarrier(ctx); ok { + for k, v := range carrierNews(row) { + response[k] = v + } + if opts.nats != nil && opts.nats.ClientTLS() { + response["nats_client_tls"] = true + } + } + return c.JSON(http.StatusOK, response) } } diff --git a/core/http/endpoints/localai/nodes_follow_test.go b/core/http/endpoints/localai/nodes_follow_test.go new file mode 100644 index 000000000000..4a25f83ac482 --- /dev/null +++ b/core/http/endpoints/localai/nodes_follow_test.go @@ -0,0 +1,280 @@ +package localai + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "time" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/pkg/natsauth" + "github.com/nats-io/nkeys" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fixedNATS is the handover of a deployment with a NATS server. +type fixedNATS struct { + url, ca string + clientTLS bool + err error +} + +func (f fixedNATS) WorkerURL(context.Context) (string, error) { return f.url, f.err } +func (f fixedNATS) CAPEM() (string, error) { return f.ca, nil } +func (f fixedNATS) ClientTLS() bool { return f.clientTLS } + +var _ = Describe("What a worker needs to follow a change of carrier", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + natsCfg natsauth.Config + tunnels *recordingDisconnector + ) + + stable := func(c cluster.Carrier) cluster.CarrierRow { + return cluster.CarrierRow{Active: c, Epoch: 5, State: cluster.StateStable} + } + preparing := func(from, to cluster.Carrier) cluster.CarrierRow { + return cluster.CarrierRow{Active: from, Epoch: 6, State: cluster.StatePrepare, Target: to} + } + handover := WithNATSHandover(fixedNATS{url: "nats://nats.example:4222", ca: "-----BEGIN CERTIFICATE-----\nCA\n-----END CERTIFICATE-----\n"}) + + BeforeEach(func() { + ctx = context.Background() + var err error + registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + tunnels = &recordingDisconnector{} + akp, err := nkeys.CreateAccount() + Expect(err).ToNot(HaveOccurred()) + seed, err := akp.Seed() + Expect(err).ToNot(HaveOccurred()) + natsCfg = natsauth.Config{AccountSeed: string(seed)} + }) + + register := func(body string, opts ...RegisterOption) map[string]any { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + opts = append(opts, WithTunnelDisconnector(tunnels)) + Expect(RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg, opts...)(e.NewContext(req, rec))).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusCreated), rec.Body.String()) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp + } + storedHash := func(name string) string { + GinkgoHelper() + node, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return node.TunnelTokenHash + } + + Describe("the registration", func() { + It("hands over the NATS address and the CA with the credential, when NATS is active", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp["nats_url"]).To(Equal("nats://nats.example:4222")) + Expect(resp["nats_ca"]).To(ContainSubstring("BEGIN CERTIFICATE")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp).ToNot(HaveKey("nats_client_tls")) + }) + + It("says that the worker needs a certificate of its own when the server asks for one", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), + WithNATSHandover(fixedNATS{url: "tls://nats.example:4222", clientTLS: true})) + Expect(resp["nats_client_tls"]).To(BeTrue()) + }) + + It("hands over nothing of NATS when none is configured, and when the read fails", func() { + reader := WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}) + Expect(register(`{"name":"a","address":"10.0.0.1:50051"}`, reader, WithNATSHandover(fixedNATS{}))).ToNot(HaveKey("nats_url")) + Expect(register(`{"name":"b","address":"10.0.0.1:50051"}`, reader, WithNATSHandover(fixedNATS{err: errors.New("down")}))).ToNot(HaveKey("nats_url")) + }) + + It("does not hand over the address of NATS to a worker that asks for the tunnel", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierNATS, cluster.CarrierTunnel)}), handover) + Expect(resp).To(HaveKey("tunnel_token")) + Expect(resp).ToNot(HaveKey("nats_url")) + Expect(resp["credential_for"]).To(Equal("tunnel")) + }) + + It("gives the credential of the target to a worker that asks for it while the change is prepared", func() { + By("a tunnel worker that follows to NATS") + first := register(`{"name":"w","address":"10.0.0.1:50051"}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierTunnel)}), handover) + held := storedHash("w") + Expect(first).To(HaveKey("tunnel_token")) + before := len(tunnels.seen()) + + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"nats"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierTunnel, cluster.CarrierNATS)}), handover) + + Expect(resp["carrier"]).To(Equal("tunnel"), "the active carrier is still the tunnel") + Expect(resp["carrier_state"]).To(Equal("prepare")) + Expect(resp["carrier_target"]).To(Equal("nats")) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp).To(HaveKey("nats_url")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(storedHash("w")).To(Equal(held), "the tunnel credential that the worker holds stays valid") + Expect(tunnels.seen()).To(HaveLen(before), "and its session is not ended") + }) + + It("mints a tunnel credential for a worker on NATS that asks for the tunnel, and leaves its NATS credential alone", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierNATS, cluster.CarrierTunnel)}), handover) + Expect(resp["carrier"]).To(Equal("nats")) + Expect(resp["credential_for"]).To(Equal("tunnel")) + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + Expect(resp).ToNot(HaveKey("nats_jwt")) + Expect(storedHash("w")).ToNot(BeEmpty()) + }) + + It("ignores a request for a carrier that is neither active, nor the target, nor draining", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`, + WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp).ToNot(HaveKey("tunnel_token"), "a node cannot get a tunnel while the cluster has none") + Expect(storedHash("w")).To(BeEmpty()) + + resp = register(`{"name":"w","address":"10.0.0.1:50051","carrier":"nonsense"}`, + WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp["credential_for"]).To(Equal("nats")) + }) + + It("gives the credential of the carrier that drains to a worker that asks for it", func() { + until := time.Now().Add(time.Minute) + row := cluster.CarrierRow{Active: cluster.CarrierTunnel, Epoch: 8, State: cluster.StateStable, Draining: cluster.CarrierNATS, DrainingUntil: &until} + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"nats"}`, WithCarrierReader(fixedCarrier{row: row}), handover) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp["carrier_draining"]).To(Equal("nats")) + Expect(resp).To(HaveKey("carrier_draining_until")) + Expect(resp).To(HaveKey("nats_jwt")) + }) + + It("keeps the address of a worker that has none to offer the tunnel while it asks for NATS", func() { + register(`{"name":"w","address":"10.0.0.1:50051","routable":true}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierTunnel)}), handover) + register(`{"name":"w","address":"10.0.0.1:50051","routable":false,"carrier":"nats"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierTunnel, cluster.CarrierNATS)}), handover) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051"), "NATS dials the address, so it must not be dropped") + }) + + It("answers a deployment with no carrier reader as it always did, whatever the worker asks", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`) + Expect(resp).ToNot(HaveKey("carrier")) + Expect(resp).ToNot(HaveKey("credential_for")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(resp).To(HaveKey("nats_jwt")) + }) + }) + + Describe("the heartbeat", func() { + var id string + + BeforeEach(func() { + id = register(`{"name":"w","address":"10.0.0.1:50051"}`)["id"].(string) + }) + + beat := func(body string, opts ...RegisterOption) (int, map[string]any) { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(HeartbeatEndpoint(registry, opts...)(c)).To(Succeed()) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return rec.Code, resp + } + + It("tells the worker the carrier, the epoch, the state and the target of a change", func() { + code, resp := beat(`{}`, WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierNATS, cluster.CarrierTunnel)})) + Expect(code).To(Equal(http.StatusOK)) + Expect(resp["message"]).To(Equal("heartbeat received")) + Expect(resp["carrier"]).To(Equal("nats")) + Expect(resp["carrier_epoch"]).To(BeEquivalentTo(6)) + Expect(resp["carrier_state"]).To(Equal("prepare")) + Expect(resp["carrier_target"]).To(Equal("tunnel")) + }) + + It("tells the worker that NATS asks for a client certificate", func() { + _, resp := beat(`{}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), WithNATSHandover(fixedNATS{url: "tls://n:4222", clientTLS: true})) + Expect(resp["nats_client_tls"]).To(BeTrue()) + _, resp = beat(`{}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp).ToNot(HaveKey("nats_client_tls")) + }) + + It("answers as it always did when there is no carrier reader, or the row cannot be read", func() { + _, resp := beat(`{}`) + Expect(resp).To(Equal(map[string]any{"message": "heartbeat received"})) + _, resp = beat(`{}`, WithCarrierReader(fixedCarrier{err: errors.New("down")})) + Expect(resp).To(Equal(map[string]any{"message": "heartbeat received"})) + }) + + It("stores what the worker reports about its carriers", func() { + code, _ := beat(`{"attached":["nats","tunnel"],"attached_epoch":6,"follow_capabilities":["nats","tunnel"],"follow_error":""}`) + Expect(code).To(Equal(http.StatusOK)) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(Equal("nats,tunnel")) + Expect(node.AttachedEpoch).To(BeEquivalentTo(6)) + Expect(node.Follow).To(Equal("nats,tunnel")) + + beat(`{"attached":["tunnel"],"attached_epoch":7,"follow_capabilities":["tunnel"],"follow_error":"no address that a frontend can dial"}`) + node, err = registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(Equal("tunnel")) + Expect(node.Follow).To(Equal("tunnel")) + Expect(node.FollowError).To(Equal("no address that a frontend can dial")) + }) + + It("stores an empty attachment as a report, because the worker still reports what it can follow", func() { + beat(`{"attached":["nats"],"follow_capabilities":["nats"]}`) + beat(`{"attached":[],"follow_capabilities":["nats","tunnel"]}`) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(BeEmpty()) + Expect(node.Follow).To(Equal("nats,tunnel")) + }) + + It("stores nothing for a worker that predates the report, and leaves an earlier report alone", func() { + beat(`{"available_ram":1024}`) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Follow).To(BeEmpty()) + Expect(node.Attached).To(BeEmpty()) + }) + + It("keeps only carriers that the cluster knows and cuts a reason that does not fit the column", func() { + beat(`{"attached":["nats","gopher","nats"],"follow_capabilities":["tunnel","nope"],"follow_error":"` + strings.Repeat("é", 200) + `"}`) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(Equal("nats")) + Expect(node.Follow).To(Equal("tunnel")) + Expect(len(node.FollowError)).To(BeNumerically("<=", 255)) + Expect(node.FollowError).To(HavePrefix("é")) + }) + + It("answers 404 for a node that is not registered and stores nothing", func() { + id = "no-such-node" + code, _ := beat(`{"attached":["nats"],"follow_capabilities":["nats"]}`) + Expect(code).To(Equal(http.StatusNotFound)) + }) + }) +}) diff --git a/core/http/routes/nodes.go b/core/http/routes/nodes.go index ecf0468f0aa7..40f35c43d7f9 100644 --- a/core/http/routes/nodes.go +++ b/core/http/routes/nodes.go @@ -46,7 +46,7 @@ func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, r node := e.Group("/api/node", readyMw, tokenAuthMw) node.POST("/register", localai.RegisterNodeEndpoint(registry, registrationToken, autoApprove, authDB, hmacSecret, natsCfg, options...)) - node.POST("/:id/heartbeat", localai.HeartbeatEndpoint(registry)) + node.POST("/:id/heartbeat", localai.HeartbeatEndpoint(registry, options...)) node.POST("/:id/drain", localai.DrainNodeEndpoint(registry, nil)) node.POST("/:id/resume", localai.ResumeNodeEndpoint(registry)) node.POST("/:id/deregister", localai.DeactivateNodeEndpoint(registry)) diff --git a/core/services/nodes/follow_report.go b/core/services/nodes/follow_report.go index 9493fab64886..81c38b116cc1 100644 --- a/core/services/nodes/follow_report.go +++ b/core/services/nodes/follow_report.go @@ -5,9 +5,11 @@ import ( "errors" "fmt" "net" + "slices" "strings" "sync" "time" + "unicode/utf8" "gorm.io/gorm" @@ -25,6 +27,32 @@ type CarrierReport struct { FollowError string } +// followErrorMax is the size of the column that holds FollowError. +const followErrorMax = 255 + +// parseCarriers keeps the names of the carriers that the cluster knows, once each. +func parseCarriers(names []string) []cluster.Carrier { + var out []cluster.Carrier + for _, n := range names { + c := cluster.Carrier(strings.TrimSpace(n)) + if (c == cluster.CarrierNATS || c == cluster.CarrierTunnel) && !slices.Contains(out, c) { + out = append(out, c) + } + } + return out +} + +// truncateRunes cuts s to at most n bytes without splitting a character. +func truncateRunes(s string, n int) string { + if len(s) <= n { + return s + } + for n > 0 && !utf8.RuneStart(s[n]) { + n-- + } + return s[:n] +} + func joinCarriers(cs []cluster.Carrier) string { out := make([]string, len(cs)) for i, c := range cs { diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index bd5b0e029051..7bdcf86c3085 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -1162,6 +1162,31 @@ type HeartbeatUpdate struct { CPULoad1 *float64 `json:"cpu_load_1,omitempty"` // WorkerIncarnation is the worker process identity. See BackendNode. WorkerIncarnation string `json:"worker_incarnation,omitempty"` + + // The carrier report. A worker that follows a change of carrier sends it in + // every heartbeat: the carriers it is attached to, the epoch of the cluster + // row it saw, the carriers it can attach to, and why it cannot attach to + // another. A worker that predates carrier switching sends none of them. + // FollowCapabilities is never empty in a report (a worker can follow the + // carrier it holds), so it tells a report from its absence. + Attached []string `json:"attached,omitempty"` + AttachedEpoch int64 `json:"attached_epoch,omitempty"` + FollowCapabilities []string `json:"follow_capabilities,omitempty"` + FollowError string `json:"follow_error,omitempty"` +} + +// CarrierReport returns the carrier report of a heartbeat, and false when the +// worker sent none. +func (u HeartbeatUpdate) CarrierReport() (CarrierReport, bool) { + if len(u.FollowCapabilities) == 0 { + return CarrierReport{}, false + } + return CarrierReport{ + Attached: parseCarriers(u.Attached), + AttachedEpoch: u.AttachedEpoch, + Follow: parseCarriers(u.FollowCapabilities), + FollowError: truncateRunes(u.FollowError, followErrorMax), + }, true } func clampCPUUsage(usage float64) float64 { From a3755e4e7b57bb8d10e9bf3322466c4f5608ed26 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:46:22 +0000 Subject: [PATCH 122/162] refactor(nodes): read the carriers of a worker from its report only The switch guessed where a worker was attached from the tunnel presence it held, because nothing reported it. Workers now report the carriers they are attached to in every heartbeat, so the guess is gone. A worker that reports its capabilities is attached to what it reports, and to nothing when it reports none: it is between two carriers and is not said to hold one it may not. A worker that reports nothing predates carrier switching, has only NATS, and is on NATS. The reader no longer takes a presence source or a reconnect grace. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 2 +- .../carrier/switch_integration_test.go | 2 +- core/services/cluster/switch.go | 2 +- core/services/nodes/follow_report.go | 54 ++++++----------- core/services/nodes/follow_report_test.go | 60 +++++++++---------- 5 files changed, 50 insertions(+), 70 deletions(-) diff --git a/core/application/distributed.go b/core/application/distributed.go index 2660d165a6da..547304f684f1 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -390,7 +390,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // While two carriers are attached, a call to a worker goes to the carrier the // worker is attached to. - workers := nodes.NewSwitchWorkers(registry, clusterReg, cluster.DefaultReconnectGrace, cfg.Distributed.StaleNodeThresholdOrDefault()) + workers := nodes.NewSwitchWorkers(registry, cfg.Distributed.StaleNodeThresholdOrDefault()) window := carrier.NewWindow( func(ctx context.Context, nodeID string) (carrier.Attachment, error) { attached, err := workers.AttachedCarriers(ctx, nodeID) diff --git a/core/services/carrier/switch_integration_test.go b/core/services/carrier/switch_integration_test.go index 50843d18fafc..c799e1bde190 100644 --- a/core/services/carrier/switch_integration_test.go +++ b/core/services/carrier/switch_integration_test.go @@ -368,7 +368,7 @@ func (rig *switchRig) startReplica(id string) *switchReplica { r.cur.Store(first) r.bus = carrier.NewBroadcaster(&r.cur) - workers := nodes.NewSwitchWorkers(rig.nodeReg, rig.clusterR, cluster.DefaultReconnectGrace, time.Hour) + workers := nodes.NewSwitchWorkers(rig.nodeReg, time.Hour) r.window = carrier.NewWindow(func(ctx context.Context, nodeID string) (carrier.Attachment, error) { cs, err := workers.AttachedCarriers(ctx, nodeID) if err != nil { diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index 944f758bf842..2198899afd01 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -62,7 +62,7 @@ func (w WorkerInfo) canFollow(target Carrier) (bool, string) { return true, "" } if !w.Reports { - return false, "the worker predates carrier switching: it reports no capabilities and stays on its carrier" + return false, "the worker predates carrier switching: it reports no capabilities and stays on NATS" } if slices.Contains(w.Follow, target) { return true, "" diff --git a/core/services/nodes/follow_report.go b/core/services/nodes/follow_report.go index 81c38b116cc1..3230f6550900 100644 --- a/core/services/nodes/follow_report.go +++ b/core/services/nodes/follow_report.go @@ -91,43 +91,32 @@ func (r *NodeRegistry) SetCarrierReport(ctx context.Context, nodeID string, rep return nil } -// PresenceReader says whether a worker holds a tunnel. *cluster.Registry is one. -type PresenceReader interface { - Presence(ctx context.Context, nodeID string, grace time.Duration) (cluster.Presence, error) -} - // SwitchWorkers answers the questions that a change of carrier asks about the // workers: which are alive, which carriers each is attached to, and which it can -// follow. A worker that reports its carriers is believed. One that does not is a -// worker that predates carrier switching: it holds a tunnel, or it is on NATS. +// follow. What it says comes from what the worker reports in its heartbeat, and +// from nothing else. A worker that reports nothing predates carrier switching: it +// has only NATS, it cannot follow, and it is on NATS. type SwitchWorkers struct { - reg *NodeRegistry - presence PresenceReader - grace time.Duration - stale time.Duration + reg *NodeRegistry + stale time.Duration } // NewSwitchWorkers returns the reader. A worker whose heartbeat is older than -// stale is not alive for this purpose; grace is the reconnect grace of the -// tunnel. -func NewSwitchWorkers(reg *NodeRegistry, presence PresenceReader, grace, stale time.Duration) *SwitchWorkers { - return &SwitchWorkers{reg: reg, presence: presence, grace: grace, stale: stale} +// stale is not alive for this purpose. +func NewSwitchWorkers(reg *NodeRegistry, stale time.Duration) *SwitchWorkers { + return &SwitchWorkers{reg: reg, stale: stale} } var _ cluster.WorkerSource = (*SwitchWorkers)(nil) -func (s *SwitchWorkers) attached(ctx context.Context, n *BackendNode) ([]cluster.Carrier, error) { - if reported := splitCarriers(n.Attached); len(reported) > 0 { - return reported, nil - } - p, err := s.presence.Presence(ctx, n.ID, s.grace) - if err != nil { - return nil, err - } - if p == cluster.PresenceConnected || p == cluster.PresenceReconnecting { - return []cluster.Carrier{cluster.CarrierTunnel}, nil +// attached returns the carriers a worker reported. A worker that reports its +// capabilities and an empty list is attached to none: it is between two +// carriers, and it is not said to be on a carrier it may not hold. +func (s *SwitchWorkers) attached(n *BackendNode) []cluster.Carrier { + if n.Follow == "" { + return []cluster.Carrier{cluster.CarrierNATS} } - return []cluster.Carrier{cluster.CarrierNATS}, nil + return splitCarriers(n.Attached) } // Workers lists the workers that are alive. @@ -142,10 +131,7 @@ func (s *SwitchWorkers) Workers(ctx context.Context) ([]cluster.WorkerInfo, erro if n.Status == StatusPending || n.Status == StatusOffline || time.Since(n.LastHeartbeat) > s.stale { continue } - attached, err := s.attached(ctx, n) - if err != nil { - return nil, fmt.Errorf("finding the carriers of node %s: %w", n.Name, err) - } + attached := s.attached(n) follow := splitCarriers(n.Follow) out = append(out, cluster.WorkerInfo{ ID: n.ID, Name: n.Name, Type: n.NodeType, @@ -161,7 +147,7 @@ func (s *SwitchWorkers) AttachedCarriers(ctx context.Context, nodeID string) ([] if err != nil { return nil, err } - return s.attached(ctx, n) + return s.attached(n), nil } // AgentsAttached reports whether an agent worker that is alive is attached to c. @@ -175,11 +161,7 @@ func (s *SwitchWorkers) AgentsAttached(ctx context.Context, c cluster.Carrier) ( if n.NodeType != NodeTypeAgent || n.Status == StatusPending || n.Status == StatusOffline || time.Since(n.LastHeartbeat) > s.stale { continue } - attached, err := s.attached(ctx, n) - if err != nil { - return false, err - } - for _, a := range attached { + for _, a := range s.attached(n) { if a == c { return true, nil } diff --git a/core/services/nodes/follow_report_test.go b/core/services/nodes/follow_report_test.go index e74915666e2e..b4fafe286a03 100644 --- a/core/services/nodes/follow_report_test.go +++ b/core/services/nodes/follow_report_test.go @@ -19,7 +19,6 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { ctx context.Context db *gorm.DB reg *NodeRegistry - presence *fakePresence workers *SwitchWorkers ) @@ -43,8 +42,7 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { var err error reg, err = NewNodeRegistry(db) Expect(err).ToNot(HaveOccurred()) - presence = &fakePresence{by: map[string]cluster.Presence{}} - workers = NewSwitchWorkers(reg, presence, cluster.DefaultReconnectGrace, 5*time.Minute) + workers = NewSwitchWorkers(reg, 5*time.Minute) }) It("lists the workers that are alive and leaves out the ones that are not", func() { @@ -62,34 +60,17 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { Expect(byName(list)).To(HaveKey("live")) }) - It("counts a worker that reports no capabilities as one that cannot follow, and says it is on NATS when it holds no tunnel", func() { + It("counts a worker that reports no capabilities as one that cannot follow, and as one that is on NATS", func() { register("old", NodeTypeBackend, "10.0.0.1:50051") list, err := workers.Workers(ctx) Expect(err).ToNot(HaveOccurred()) w := byName(list)["old"] Expect(w.Reports).To(BeFalse()) - Expect(w.Attached).To(Equal([]cluster.Carrier{cluster.CarrierNATS})) + Expect(w.Attached).To(Equal([]cluster.Carrier{cluster.CarrierNATS}), "a worker that predates carrier switching has no other carrier") }) - It("says a worker that holds a tunnel is attached to the tunnel", func() { - n := register("tun", NodeTypeBackend, "") - presence.by[n.ID] = cluster.PresenceConnected - list, err := workers.Workers(ctx) - Expect(err).ToNot(HaveOccurred()) - Expect(byName(list)["tun"].Attached).To(Equal([]cluster.Carrier{cluster.CarrierTunnel})) - }) - - It("says a worker that was attached to the tunnel a moment ago is still attached to it", func() { - n := register("blink", NodeTypeBackend, "") - presence.by[n.ID] = cluster.PresenceReconnecting - list, err := workers.Workers(ctx) - Expect(err).ToNot(HaveOccurred()) - Expect(byName(list)["blink"].Attached).To(Equal([]cluster.Carrier{cluster.CarrierTunnel})) - }) - - It("takes what a worker reports over what can be inferred", func() { + It("takes the carriers a worker reports, and nothing else", func() { n := register("both", NodeTypeBackend, "10.0.0.1:50051") - presence.by[n.ID] = cluster.PresenceConnected Expect(reg.SetCarrierReport(ctx, n.ID, CarrierReport{ Attached: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, AttachedEpoch: 7, Follow: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, FollowError: "", @@ -103,6 +84,30 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { Expect(w.Follow).To(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) }) + It("says that a worker which reports a tunnel is on the tunnel, and one which reports NATS is on NATS", func() { + a := register("tun", NodeTypeBackend, "") + b := register("nats", NodeTypeBackend, "10.0.0.1:50051") + Expect(reg.SetCarrierReport(ctx, a.ID, CarrierReport{Attached: []cluster.Carrier{cluster.CarrierTunnel}, Follow: []cluster.Carrier{cluster.CarrierTunnel}})).To(Succeed()) + Expect(reg.SetCarrierReport(ctx, b.ID, CarrierReport{Attached: []cluster.Carrier{cluster.CarrierNATS}, Follow: []cluster.Carrier{cluster.CarrierNATS}})).To(Succeed()) + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(byName(list)["tun"].Attached).To(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + Expect(byName(list)["nats"].Attached).To(Equal([]cluster.Carrier{cluster.CarrierNATS})) + }) + + It("does not guess a carrier for a worker that reports it is attached to none", func() { + n := register("between", NodeTypeBackend, "") + Expect(reg.SetCarrierReport(ctx, n.ID, CarrierReport{Attached: nil, Follow: []cluster.Carrier{cluster.CarrierTunnel}})).To(Succeed()) + list, err := workers.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + w := byName(list)["between"] + Expect(w.Reports).To(BeTrue()) + Expect(w.Attached).To(BeEmpty()) + a, err := workers.AttachedCarriers(ctx, n.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(a).To(BeEmpty()) + }) + It("carries the reason a worker gives for not following", func() { n := register("tunnel-only", NodeTypeBackend, "") Expect(reg.SetCarrierReport(ctx, n.ID, CarrierReport{ @@ -127,7 +132,7 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { It("says whether an agent worker is attached to a carrier", func() { register("backend", NodeTypeBackend, "10.0.0.1:50051") agent := register("agent", NodeTypeAgent, "") - presence.by[agent.ID] = cluster.PresenceConnected + Expect(reg.SetCarrierReport(ctx, agent.ID, CarrierReport{Attached: []cluster.Carrier{cluster.CarrierTunnel}, Follow: []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}})).To(Succeed()) onTunnel, err := workers.AgentsAttached(ctx, cluster.CarrierTunnel) Expect(err).ToNot(HaveOccurred()) @@ -136,13 +141,6 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { Expect(err).ToNot(HaveOccurred()) Expect(onNATS).To(BeFalse(), "the backend worker is on NATS, and it is not an agent worker") }) - - It("returns the error of the presence read and not a guess", func() { - n := register("w", NodeTypeBackend, "") - presence.err = errors.New("the database is down") - _, err := workers.AttachedCarriers(ctx, n.ID) - Expect(err).To(HaveOccurred()) - }) }) var _ = Describe("The direct client factory for a worker that holds a tunnel", func() { From 4c28ab3a112335e183d6efe39ad288e286d19b6b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:47:16 +0000 Subject: [PATCH 123/162] feat(workerregistry): ask for the credential of one carrier and read the heartbeat answer A worker that follows a change of carrier waits for the credential of the carrier it attaches to, which is not the active one while the change is prepared. A credential manager can be made for one carrier. It refuses an answer that carries the credential of another carrier, as a frontend that does not allow the request would send, and tries again. The registration answer carries the state of a change, the address of NATS, its CA and the client certificate flag. The heartbeat client can return the answer of the frontend, which names the carrier, the epoch, and the target of a change. A frontend that predates carriers answers without them. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/workerregistry/client.go | 66 +++++++++++++-- core/cli/workerregistry/credentials.go | 41 ++++++++- core/cli/workerregistry/follow_test.go | 110 +++++++++++++++++++++++++ 3 files changed, 205 insertions(+), 12 deletions(-) create mode 100644 core/cli/workerregistry/follow_test.go diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index feb5e882b3d7..e5409246aec3 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -9,6 +9,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "strings" "sync" @@ -71,6 +72,31 @@ type RegisterResponse struct { // TunnelToken is the own credential of the node for the tunnel. It is sent // once and only when the tunnel is the active carrier. TunnelToken string `json:"tunnel_token,omitempty"` + // CarrierState, CarrierTarget and CarrierDraining say whether a change of + // carrier is under way. CredentialFor is the carrier that the credential in + // the answer is for: the one the worker asked for, when the frontend allowed + // it, else the active one. + CarrierState string `json:"carrier_state,omitempty"` + CarrierTarget string `json:"carrier_target,omitempty"` + CarrierDraining string `json:"carrier_draining,omitempty"` + CredentialFor string `json:"credential_for,omitempty"` + // NatsURL is the address workers use for NATS, NatsCA the CA of the server as + // PEM, and NatsClientTLS says the server asks for a client certificate, which + // the frontend cannot hand over. + NatsURL string `json:"nats_url,omitempty"` + NatsCA string `json:"nats_ca,omitempty"` + NatsClientTLS bool `json:"nats_client_tls,omitempty"` +} + +// HeartbeatReply is the answer to a heartbeat. The carrier fields are empty when +// the frontend predates carriers. +type HeartbeatReply struct { + Carrier string `json:"carrier,omitempty"` + CarrierEpoch int64 `json:"carrier_epoch,omitempty"` + CarrierState string `json:"carrier_state,omitempty"` + CarrierTarget string `json:"carrier_target,omitempty"` + CarrierDraining string `json:"carrier_draining,omitempty"` + NatsClientTLS bool `json:"nats_client_tls,omitempty"` } // RegisterFull sends a single registration request and returns the full @@ -143,24 +169,46 @@ func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map return nil, err } -// Heartbeat sends a single heartbeat POST with the given body. +// Heartbeat sends a single heartbeat POST with the given body. It fails only when +// the request cannot be made. func (c *RegistrationClient) Heartbeat(ctx context.Context, nodeID string, body map[string]any) error { + resp, err := c.post(ctx, nodeID, body) + if err != nil { + return err + } + return resp.Body.Close() +} + +// HeartbeatFull sends a heartbeat and returns the answer, which tells the worker +// which carrier the cluster uses. It fails when the frontend refuses the +// heartbeat, for example for a node it does not know. +func (c *RegistrationClient) HeartbeatFull(ctx context.Context, nodeID string, body map[string]any) (*HeartbeatReply, error) { + resp, err := c.post(ctx, nodeID, body) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("heartbeat refused with status %d", resp.StatusCode) + } + var reply HeartbeatReply + // An answer that is not JSON is the answer of a frontend that says nothing + // about carriers. + _ = json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&reply) + return &reply, nil +} + +func (c *RegistrationClient) post(ctx context.Context, nodeID string, body map[string]any) (*http.Response, error) { jsonBody, _ := json.Marshal(body) url := c.baseURL() + "/api/node/" + nodeID + "/heartbeat" req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(jsonBody)) if err != nil { - return fmt.Errorf("creating heartbeat request: %w", err) + return nil, fmt.Errorf("creating heartbeat request: %w", err) } req.Header.Set("Content-Type", "application/json") c.setAuth(req) - - resp, err := c.httpClient().Do(req) - if err != nil { - return err - } - defer resp.Body.Close() - return nil + return c.httpClient().Do(req) } // HeartbeatLoop runs heartbeats at the given interval until ctx is cancelled. diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 2b052f7df4a1..d39a6cd0e449 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -43,6 +43,12 @@ type RegisterFunc func(ctx context.Context) (*RegisterResponse, error) type CredentialManager struct { register RegisterFunc requireCreds bool // block until credentials are present (frontend minting in use) + // carrier is the carrier whose credential this manager waits for. Empty + // means the one the frontend names as active, which is what a worker wants + // when it starts. A worker that follows a change asks for the carrier it + // attaches to, and that is not the active one yet while the change is + // prepared. + carrier string // Tunables; defaults set by NewCredentialManager, overridable in tests. initialBackoff time.Duration @@ -59,6 +65,29 @@ type CredentialManager struct { tunnelToken string } +// NewCredentialManagerFor is NewCredentialManager for a manager that waits for the +// credential of one carrier, whichever the cluster has active. register must ask +// the frontend for that carrier. +func NewCredentialManagerFor(register RegisterFunc, requireCreds bool, carrier string) *CredentialManager { + m := NewCredentialManager(register, requireCreds) + m.carrier = carrier + return m +} + +// credentialFor is the carrier that an answer carries a credential for. +func (m *CredentialManager) credentialFor(res *RegisterResponse) string { + if m.carrier != "" { + return m.carrier + } + return res.Carrier +} + +// wrongCarrier is true when the manager asked for a carrier and the frontend +// answered with the credential of another. +func (m *CredentialManager) wrongCarrier(res *RegisterResponse) bool { + return m.carrier != "" && res.CredentialFor != "" && res.CredentialFor != m.carrier +} + // NewCredentialManager builds a manager over register. When requireCreds is // true, Acquire blocks until the node is approved and credentials are minted. func NewCredentialManager(register RegisterFunc, requireCreds bool) *CredentialManager { @@ -153,10 +182,13 @@ func (m *CredentialManager) Acquire(ctx context.Context) (*RegisterResponse, err case err != nil: lastReason = err xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) - case res.Carrier == carrierTunnel && res.TunnelToken == "": + case m.wrongCarrier(res): + lastReason = fmt.Errorf("asked for the credential of the %s carrier and got the one of the %s carrier", m.carrier, res.CredentialFor) + xlog.Info("The frontend does not hand out the credential of that carrier yet; waiting", "asked", m.carrier, "got", res.CredentialFor, "attempt", attempt, "next_retry", backoff) + case m.credentialFor(res) == carrierTunnel && res.TunnelToken == "": lastReason = fmt.Errorf("node %s registered but the tunnel credential was not minted", res.ID) xlog.Info("Node registered but the tunnel credential is not minted yet; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff) - case res.Carrier == carrierTunnel: + case m.credentialFor(res) == carrierTunnel: m.store(res) return res, nil case !m.requireCreds: @@ -193,7 +225,10 @@ func (m *CredentialManager) Reregister(ctx context.Context) error { if err != nil { return err } - if res.Carrier == carrierTunnel && res.TunnelToken == "" { + if m.wrongCarrier(res) { + return fmt.Errorf("asked for the credential of the %s carrier and got the one of the %s carrier", m.carrier, res.CredentialFor) + } + if m.credentialFor(res) == carrierTunnel && res.TunnelToken == "" { return fmt.Errorf("node %s registered but the tunnel credential was not minted", res.ID) } m.store(res) diff --git a/core/cli/workerregistry/follow_test.go b/core/cli/workerregistry/follow_test.go new file mode 100644 index 000000000000..96674e1a0653 --- /dev/null +++ b/core/cli/workerregistry/follow_test.go @@ -0,0 +1,110 @@ +package workerregistry + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("A credential for the carrier a worker follows to", func() { + It("waits for the tunnel token when it asks for the tunnel, although the active carrier is NATS", func() { + f := &fakeRegister{steps: []step{ + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel"}}, + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel", TunnelToken: "tok"}}, + }} + m := NewCredentialManagerFor(f.fn(), false, "tunnel") + m.initialBackoff = 1 + res, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(res.TunnelToken).To(Equal("tok")) + Expect(m.TunnelToken()).To(Equal("tok")) + Expect(f.count()).To(Equal(2)) + }) + + It("takes the NATS credential while the tunnel is the active carrier", func() { + f := &fakeRegister{steps: []step{{res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "tunnel", CredentialFor: "nats", NatsJWT: "j", NatsUserSeed: "s", NatsURL: "nats://x:4222"}}}} + m := NewCredentialManagerFor(f.fn(), true, "nats") + res, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(res.NatsURL).To(Equal("nats://x:4222")) + jwt, seed := m.Current() + Expect(jwt).To(Equal("j")) + Expect(seed).To(Equal("s")) + }) + + It("does not take the credential of another carrier for the one it asked for", func() { + f := &fakeRegister{steps: []step{{res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "nats", NatsJWT: "j", NatsUserSeed: "s"}}}} + m := NewCredentialManagerFor(f.fn(), false, "tunnel") + m.maxAttempts = 2 + m.initialBackoff = 1 + _, err := m.Acquire(context.Background()) + Expect(err).To(MatchError(ContainSubstring("tunnel"))) + Expect(m.TunnelToken()).To(BeEmpty()) + }) + + It("asks again for the same carrier when it registers again", func() { + f := &fakeRegister{steps: []step{ + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel", TunnelToken: "a"}}, + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel", TunnelToken: "b"}}, + }} + m := NewCredentialManagerFor(f.fn(), false, "tunnel") + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(m.Reregister(context.Background())).To(Succeed()) + Expect(m.TunnelToken()).To(Equal("b")) + }) +}) + +var _ = Describe("The heartbeat client", func() { + It("returns what the frontend says about the carrier", func() { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + Expect(r.URL.Path).To(Equal("/api/node/n1/heartbeat")) + var body map[string]any + Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed()) + Expect(body).To(HaveKey("follow_capabilities")) + _ = json.NewEncoder(w).Encode(map[string]any{"message": "heartbeat received", "carrier": "nats", "carrier_epoch": 4, "carrier_state": "prepare", "carrier_target": "tunnel", "nats_client_tls": true}) + })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + reply, err := c.HeartbeatFull(context.Background(), "n1", map[string]any{"follow_capabilities": []string{"nats"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Carrier).To(Equal("nats")) + Expect(reply.CarrierEpoch).To(BeEquivalentTo(4)) + Expect(reply.CarrierState).To(Equal("prepare")) + Expect(reply.CarrierTarget).To(Equal("tunnel")) + Expect(reply.NatsClientTLS).To(BeTrue()) + }) + + It("returns an empty answer from a frontend that predates carriers", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]string{"message": "heartbeat received"}) + })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + reply, err := c.HeartbeatFull(context.Background(), "n1", nil) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Carrier).To(BeEmpty()) + }) + + It("reports a refusal and a node the frontend does not know", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNotFound) })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + _, err := c.HeartbeatFull(context.Background(), "n1", nil) + Expect(err).To(HaveOccurred()) + }) + + It("keeps the heartbeat that returns only an error working as before", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + Expect(c.Heartbeat(context.Background(), "n1", nil)).To(Succeed()) + }) +}) From 468ffdfd93b7994b8bbced32730d07241b7c922e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:52:29 +0000 Subject: [PATCH 124/162] feat(worker): follow the carrier that the cluster names A worker used to attach to one carrier at start and keep it. It now reads the answer of every heartbeat, which reaches the frontend whichever carrier is active, and follows the carrier the frontend names. While a change is prepared it attaches to the target as well: it registers again for the credential of that carrier, waits a random time of up to ten seconds so that a fleet does not connect at once, and connects with a wait that grows from one to thirty seconds when the carrier refuses. It keeps the previous carrier open through the drain, tells the frontends where it is attached in every heartbeat, and closes a carrier only when the cluster released it, the carrier it moved to is connected, and no request runs on it. A backend process is not touched. A change that is aborted closes the target again. A worker that cannot follow keeps the carrier it has, says why in its heartbeat and tries again. It cannot follow to NATS without an address that the frontends can dial, or with a server that asks for a client certificate it does not have. A worker that boots on NATS counts as routable, because its backends listen on every interface. A frontend that predates carriers leaves the worker as it was. The NATS side takes the address, the CA and the credential from the frontend. A local address or CA still wins, so the flag keeps its meaning. The tunnel is attached and released through the same path at start and later, and the control plane of the HTTP server is empty whenever the worker holds no tunnel. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/carrier_test.go | 19 +- core/services/worker/config.go | 28 + core/services/worker/control_http.go | 3 + core/services/worker/control_nats.go | 5 + core/services/worker/control_server.go | 15 + core/services/worker/follow.go | 10 + core/services/worker/follow_attach.go | 164 +++++ core/services/worker/follow_carrier.go | 648 ++++++++++++++++++++ core/services/worker/follow_carrier_test.go | 308 ++++++++++ core/services/worker/nats_connect.go | 34 - core/services/worker/supervisor.go | 5 - core/services/worker/worker.go | 284 ++++----- 12 files changed, 1309 insertions(+), 214 deletions(-) create mode 100644 core/services/worker/follow_attach.go create mode 100644 core/services/worker/follow_carrier.go create mode 100644 core/services/worker/follow_carrier_test.go diff --git a/core/services/worker/carrier_test.go b/core/services/worker/carrier_test.go index b9245328bcdf..961e11a31275 100644 --- a/core/services/worker/carrier_test.go +++ b/core/services/worker/carrier_test.go @@ -8,7 +8,7 @@ import ( var _ = Describe("The carrier that a worker attaches to", func() { DescribeTable("follows the answer of the registration", func(named, natsURL string, wantTunnel bool, wantErr string) { - onTunnel, err := carrierOf(named, &Config{NatsURL: natsURL}) + onTunnel, err := carrierOf(named, "", &Config{NatsURL: natsURL}) if wantErr != "" { Expect(err).To(MatchError(ContainSubstring(wantErr))) return @@ -22,6 +22,23 @@ var _ = Describe("The carrier that a worker attaches to", func() { Entry("no name (a frontend before carriers), with a URL", "", "nats://n:4222", false, ""), Entry("NATS, with no URL", "nats", "", false, "no NATS URL"), Entry("no name, with no URL", "", "", false, "does not name a carrier"), + Entry("NATS, with no URL of its own and none from the frontend", "nats", "", false, "handed over none"), Entry("a carrier from the future", "carrier-pigeon", "nats://n:4222", false, "does not know"), ) + + DescribeTable("takes the address of NATS from the frontend when it has none", + func(local, handed string, wantErr bool) { + onTunnel, err := carrierOf("nats", handed, &Config{NatsURL: local}) + if wantErr { + Expect(err).To(HaveOccurred()) + return + } + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(BeFalse()) + }, + Entry("only the frontend gives one", "", "nats://frontend-gave:4222", false), + Entry("only the worker has one", "nats://local:4222", "", false), + Entry("both, and the local one wins as an override", "nats://local:4222", "nats://frontend-gave:4222", false), + Entry("neither", "", "", true), + ) }) diff --git a/core/services/worker/config.go b/core/services/worker/config.go index 74f19c0e906e..150a09542df8 100644 --- a/core/services/worker/config.go +++ b/core/services/worker/config.go @@ -1,5 +1,11 @@ package worker +import ( + "time" + + "github.com/mudler/xlog" +) + // Config is the configuration for the distributed agent worker. // // Field tags are kong/kong-env metadata read by core/cli/worker.go's WorkerCMD, @@ -78,6 +84,10 @@ type Config struct { // enforces it against the raw VRAM this worker reports. Empty = no cap. VRAMBudget string `env:"LOCALAI_VRAM_BUDGET" help:"Cap VRAM used for model allocation on this worker node, as a percentage (e.g. 80%) or absolute amount (e.g. 12GB)." group:"registration"` + // FollowMaxDelay bounds the random wait before the worker attaches to a carrier + // that the cluster changed to, so that a fleet does not connect at once. + FollowMaxDelay string `env:"LOCALAI_FOLLOW_MAX_DELAY" default:"10s" help:"Longest random wait before this worker attaches to the carrier that the cluster changed to" group:"registration" hidden:""` + // NATS. The URL is needed only when the cluster runs on NATS. NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed when the cluster runs on NATS; a cluster on the tunnel does not use it" group:"distributed"` NatsJWT string `env:"LOCALAI_NATS_JWT" help:"NATS user JWT override (normally from registration nats_jwt)" group:"distributed"` @@ -106,3 +116,21 @@ func (c Config) NatsAuthRequired() bool { func (c Config) RegistrationAuthRequired() bool { return c.RegistrationRequireAuth || c.DistributedRequireAuth } + +// followMaxDelay returns the longest random wait before the worker follows a +// change of carrier. A value that is not a duration is the default. +func (c Config) followMaxDelay() time.Duration { + if c.FollowMaxDelay == "" { + return DefaultFollowMaxDelay + } + d, err := time.ParseDuration(c.FollowMaxDelay) + if err != nil { + xlog.Warn("invalid follow delay, using the default", "input", c.FollowMaxDelay, "error", err) + return DefaultFollowMaxDelay + } + if d == 0 { + // Zero is "no wait" here, and a zero in the options means the default. + return -1 + } + return d +} diff --git a/core/services/worker/control_http.go b/core/services/worker/control_http.go index 0efddcb064e0..f64cf5e2b3f7 100644 --- a/core/services/worker/control_http.go +++ b/core/services/worker/control_http.go @@ -28,6 +28,7 @@ import ( // not serve, 405 for a method other than POST. The frontend maps each of them // to a different condition, so a handler must not use them for its own result. type httpControlServer struct { + inflight mux *http.ServeMux mu sync.Mutex @@ -79,6 +80,7 @@ func (s *httpControlServer) handle(v controlVerb, h controlHandler) error { return err } s.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + defer s.enter()() body, ok := workerctl.ReadRequestBody(w, r) if !ok { return @@ -113,6 +115,7 @@ func (s *httpControlServer) handleWithProgress(v controlVerb, h progressControlH return err } s.mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + defer s.enter()() body, ok := workerctl.ReadRequestBody(w, r) if !ok { return diff --git a/core/services/worker/control_nats.go b/core/services/worker/control_nats.go index b9e8dc7b2abc..007c503eb768 100644 --- a/core/services/worker/control_nats.go +++ b/core/services/worker/control_nats.go @@ -10,6 +10,7 @@ import ( // natsControlServer serves control verbs on this node's NATS subjects. type natsControlServer struct { + inflight bus messaging.MessagingClient nodeID string } @@ -67,6 +68,8 @@ func (n *natsControlServer) handle(v controlVerb, h controlHandler) error { return fmt.Errorf("serving %s: %w", v, err) } if _, err := n.bus.SubscribeReply(subject, func(data []byte, reply func([]byte)) { + done := n.enter() + defer done() if r, _ := h(context.Background(), data); r != nil { replyJSON(reply, r) } @@ -89,7 +92,9 @@ func (n *natsControlServer) handleWithProgress(v controlVerb, h progressControlH _ = n.bus.Publish(messaging.SubjectNodeBackendInstallProgress(n.nodeID, ev.OpID), ev) } if _, err := n.bus.SubscribeReply(subject, func(data []byte, reply func([]byte)) { + done := n.enter() go func() { + defer done() if r, _ := h(context.Background(), data, progress); r != nil { replyJSON(reply, r) } diff --git a/core/services/worker/control_server.go b/core/services/worker/control_server.go index a1b068102027..aabc07fce565 100644 --- a/core/services/worker/control_server.go +++ b/core/services/worker/control_server.go @@ -3,6 +3,7 @@ package worker import ( "context" "encoding/json" + "sync/atomic" "github.com/mudler/LocalAI/core/services/workerctl" ) @@ -111,3 +112,17 @@ func refuseNever[Reply any](error) Reply { var reply Reply return reply } + +// inflight counts the requests that a control server is answering. A worker does +// not close a carrier while one runs on it, because the answer goes back on the +// carrier the request came on. +type inflight struct{ n atomic.Int32 } + +// enter marks a request as running and returns the function that ends it. +func (c *inflight) enter() func() { + c.n.Add(1) + return func() { c.n.Add(-1) } +} + +// InFlight is the number of requests that run now. +func (c *inflight) InFlight() int { return int(c.n.Load()) } diff --git a/core/services/worker/follow.go b/core/services/worker/follow.go index f4eabd3f73fa..3f4aef714581 100644 --- a/core/services/worker/follow.go +++ b/core/services/worker/follow.go @@ -70,6 +70,16 @@ func newTunnelFollower(o tunnelFollowerOptions) *tunnelFollower { } } +// UseCredentials sets where the tunnel reads its credential, and how it gets a +// new one. A credential is new after every registration, so the follower is given +// it again for each attach. +func (f *tunnelFollower) UseCredentials(token func() string, reauthorize func(ctx context.Context) error) { + f.mu.Lock() + defer f.mu.Unlock() + f.token = token + f.reauthorize = reauthorize +} + // CanAttach reports whether this worker has what a tunnel needs: a frontend URL // that can be turned into the endpoint of the tunnel. A worker that cannot attach // must not mount a control plane at all. diff --git a/core/services/worker/follow_attach.go b/core/services/worker/follow_attach.go new file mode 100644 index 000000000000..39fada8e4dd7 --- /dev/null +++ b/core/services/worker/follow_attach.go @@ -0,0 +1,164 @@ +package worker + +import ( + "cmp" + "context" + "errors" + "sync/atomic" + + "github.com/mudler/LocalAI/core/cli/workerregistry" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/xlog" +) + +// natsAttachment is a worker's connection to NATS with its control verbs served +// on the subjects of the node. +type natsAttachment struct { + client *messaging.Client + control *natsControlServer + cancel context.CancelFunc + failed atomic.Bool +} + +func (a *natsAttachment) Connected() bool { return a.client.IsConnected() } +func (a *natsAttachment) InFlight() int { return a.control.InFlight() } +func (a *natsAttachment) Failed() bool { return a.failed.Load() } +func (a *natsAttachment) Close() error { + a.cancel() + a.client.Close() + return nil +} + +// NATSLocal is what an operator can set on a worker about NATS. It overrides what +// the frontend hands over. +type NATSLocal struct { + URL string + JWT string + Seed string + Required bool + TLS messaging.TLSFiles +} + +// natsTLSFor returns the TLS files to use: the files of the operator, and the CA +// that the frontend handed over when the operator gave none. +func natsTLSFor(local messaging.TLSFiles, h Handover) messaging.TLSFiles { + tls := local + if tls.CA == "" && h.NATSCAPEM != "" { + tls.CAPEM = []byte(h.NATSCAPEM) + } + return tls +} + +// ConnectNATS connects to NATS as a worker. The address and the CA of the +// operator win over the ones the frontend handed over. The credential comes from +// the environment when it is set there, and otherwise from the manager, which +// keeps it fresh. onFatal is called when a refresh fails for good. +func ConnectNATS(ctx context.Context, local NATSLocal, h Handover, onFatal func()) (*messaging.Client, error) { + url := cmp.Or(local.URL, h.NATSURL) + if url == "" { + return nil, errors.New("no NATS URL: set LOCALAI_NATS_URL, or set the NATS address in the cluster settings") + } + tls := natsTLSFor(local.TLS, h) + if local.JWT != "" || local.Seed != "" { + return connectNATS(url, local.JWT, local.Seed, "", "", local.Required, tls) + } + var opts []messaging.Option + if h.Creds != nil && h.Creds.HasCredentials() { + opts = append(opts, messaging.WithUserJWTProvider(h.Creds.Provider())) + } else if local.Required { + return nil, errors.New("NATS JWT+seed required: set LOCALAI_NATS_JWT/LOCALAI_NATS_USER_SEED or enable frontend minting") + } + if tls.Enabled() { + opts = append(opts, messaging.WithTLS(tls)) + } + client, err := messaging.New(url, opts...) + if err == nil && h.Creds != nil && h.Creds.HasCredentials() { + go func() { + if err := h.Creds.RefreshLoop(ctx); err != nil { + xlog.Error("NATS credential refresh permanently failed", "error", err) + onFatal() + } + }() + } + return client, err +} + +// natsAttacher attaches a backend worker to NATS and serves the verbs on it. +func natsAttacher(local NATSLocal, nodeID string, serve func(controlServer) error) Attacher { + return func(ctx context.Context, h Handover) (Attachment, error) { + actx, cancel := context.WithCancel(ctx) + a := &natsAttachment{cancel: cancel} + client, err := ConnectNATS(actx, local, h, func() { a.failed.Store(true) }) + if err != nil { + cancel() + return nil, err + } + a.client = client + a.control = newNATSControlServer(client, nodeID) + if err := serve(a.control); err != nil { + _ = a.Close() + return nil, err + } + return a, nil + } +} + +// tunnelAttachment is a worker's tunnel with its control plane mounted behind it. +type tunnelAttachment struct{ f *tunnelFollower } + +func (a tunnelAttachment) Connected() bool { + a.f.mu.Lock() + t := a.f.tunnel + a.f.mu.Unlock() + return t.Connected() +} +func (a tunnelAttachment) InFlight() int { return a.f.control.InFlight() } +func (a tunnelAttachment) Close() error { return a.f.Detach() } + +// tunnelAttacher attaches a backend worker to a tunnel. +func tunnelAttacher(f *tunnelFollower) Attacher { + return func(ctx context.Context, h Handover) (Attachment, error) { + if h.Creds == nil { + return nil, ErrNoCredentials + } + f.UseCredentials(h.Creds.TunnelToken, h.Creds.Reregister) + if err := f.Attach(ctx); err != nil { + return nil, err + } + return tunnelAttachment{f}, nil + } +} + +// NATSReasons says why a worker cannot attach to NATS, or returns an empty +// string. routable is true when the frontends can dial the worker, which is how +// they reach it on NATS. +func NATSReasons(routable bool, local NATSLocal) func(CarrierView) string { + return func(v CarrierView) string { + switch { + case !routable: + return "this worker has no address that the frontends can dial: start it with --addr to let it follow a change to NATS" + case v.NATSClientTLS && (local.TLS.Cert == "" || local.TLS.Key == ""): + return "the NATS server asks for a client certificate, and this worker has none: set LOCALAI_NATS_TLS_CERT and LOCALAI_NATS_TLS_KEY" + } + return "" + } +} + +// registerFor builds the function that registers a worker for one carrier. The +// body is rebuilt for each registration, because it carries the load of the host. +func registerFor(client *workerregistry.RegistrationClient, body func() map[string]any, c cluster.Carrier) workerregistry.RegisterFunc { + return func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + b := body() + b["carrier"] = string(c) + return client.RegisterFull(ctx, b) + } +} + +// beatOf turns the answer to a registration into the news of a heartbeat. +func beatOf(res *workerregistry.RegisterResponse) *Beat { + return &Beat{ + Carrier: res.Carrier, CarrierEpoch: res.CarrierEpoch, CarrierState: res.CarrierState, + CarrierTarget: res.CarrierTarget, CarrierDraining: res.CarrierDraining, NatsClientTLS: res.NatsClientTLS, + } +} diff --git a/core/services/worker/follow_carrier.go b/core/services/worker/follow_carrier.go new file mode 100644 index 000000000000..47e569e896c6 --- /dev/null +++ b/core/services/worker/follow_carrier.go @@ -0,0 +1,648 @@ +package worker + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "slices" + "strings" + "sync" + "time" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/cli/workerregistry" + "github.com/mudler/LocalAI/core/services/cluster" +) + +// Timings of a worker that follows a change of carrier. +const ( + // DefaultFollowMaxDelay bounds the random wait before a worker attaches to a + // carrier it follows. Workers learn of a change within a few seconds of each + // other, and without it they would all connect in the same instant. + DefaultFollowMaxDelay = 10 * time.Second + // followBackoffMin and followBackoffMax bound the wait between two attempts + // to attach to a carrier. + followBackoffMin = time.Second + followBackoffMax = 30 * time.Second + // followFastBeat is the interval of the heartbeat while a change is under + // way or a carrier is being released, so that the news moves faster than at + // the usual pace of ten seconds. + followFastBeat = 2 * time.Second + // followDrainGrace is how long a worker keeps a carrier open after the + // cluster released it, when a request is still running on it. + followDrainGrace = 2 * time.Minute + // followCredentialWait bounds the wait for the credential of a carrier. + followCredentialWait = time.Minute + // followRecheck is how often a worker that cannot follow looks again. + followRecheck = 30 * time.Second +) + +// Attachment is a worker's hold on one carrier: a NATS connection with the verbs +// served on it, or a tunnel with the control plane mounted behind it. +type Attachment interface { + // Connected reports whether the carrier can be used now. It is false while the + // attachment connects and while it connects again. + Connected() bool + // InFlight is the number of requests that run on the carrier. The carrier is + // not closed while one runs. + InFlight() int + // Close ends the hold. The process of a backend is not touched. + Close() error +} + +// failable is implemented by an attachment that can end by itself for good, for +// example when its credential cannot be renewed. The follower then attaches +// again. +type failable interface{ Failed() bool } + +// Handover is what the frontend gave a worker for one carrier. +type Handover struct { + Carrier cluster.Carrier + // Creds holds the credential of the carrier. It keeps the credential fresh. + Creds *workerregistry.CredentialManager + // The address of NATS, its CA as PEM, and whether the server asks for a + // client certificate. Set for NATS only. + NATSURL string + NATSCAPEM string + NATSClientTLS bool +} + +// Attacher connects a worker to one carrier. It returns once the attachment +// exists. It does not wait for the carrier to answer. +type Attacher func(ctx context.Context, h Handover) (Attachment, error) + +// CarrierView is what the frontend last said about the carrier of the cluster. +type CarrierView struct { + Known bool + Active cluster.Carrier + Epoch int64 + State cluster.State + Target cluster.Carrier + Draining cluster.Carrier + NATSClientTLS bool +} + +// Beat is the answer to a heartbeat or to a registration. +type Beat = workerregistry.HeartbeatReply + +// FollowerOptions configures a Follower. +type FollowerOptions struct { + // Attachers connect to each carrier the worker can hold. + Attachers map[cluster.Carrier]Attacher + // Credentials returns the manager that registers the worker for one carrier + // and keeps its credential. It asks the frontend for that carrier. + Credentials func(c cluster.Carrier) *workerregistry.CredentialManager + // Heartbeat sends a heartbeat and returns the answer. + Heartbeat func(ctx context.Context, body map[string]any) (*Beat, error) + // Body returns the usual content of a heartbeat. May be nil. + Body func() map[string]any + // Cannot says, for each carrier, why this worker cannot attach to it, or + // returns an empty string. A carrier with no entry can be attached to. + Cannot map[cluster.Carrier]func(CarrierView) string + // Interval is the usual interval of the heartbeat. + Interval time.Duration + // MaxDelay bounds the random wait before an attach. Zero uses + // DefaultFollowMaxDelay. A negative value means no wait. + MaxDelay time.Duration +} + +// Follower keeps a worker on the carrier of the cluster. +// +// The worker never decides which carrier is active. It reads the answer of every +// heartbeat, attaches to the carrier that the frontend names (and, while a change +// is prepared, to its target), keeps the previous carrier open while the cluster +// drains it, and closes it when the cluster releases it and nothing runs on it. +// A backend process keeps serving all through: only the channel of control moves. +// A worker that cannot follow stays on the carrier it has, says why in its +// heartbeat, and tries again. +// +// What a worker reports in every heartbeat is the single source for the +// frontends of where the worker is attached. +type Follower struct { + o FollowerOptions + + mu sync.Mutex + view CarrierView + attached map[cluster.Carrier]Attachment + errs map[cluster.Carrier]string + tries map[cluster.Carrier]int + retryAt map[cluster.Carrier]time.Time + delayed map[cluster.Carrier]int64 + relFrom map[cluster.Carrier]time.Time + + kick chan struct{} + beat sync.Mutex + root context.Context +} + +// NewFollower returns a follower. Run starts it. +func NewFollower(o FollowerOptions) *Follower { + if o.Interval <= 0 { + o.Interval = 10 * time.Second + } + if o.MaxDelay == 0 { + o.MaxDelay = DefaultFollowMaxDelay + } + return &Follower{ + o: o, + attached: map[cluster.Carrier]Attachment{}, + errs: map[cluster.Carrier]string{}, + tries: map[cluster.Carrier]int{}, + retryAt: map[cluster.Carrier]time.Time{}, + delayed: map[cluster.Carrier]int64{}, + relFrom: map[cluster.Carrier]time.Time{}, + kick: make(chan struct{}, 1), + } +} + +// Adopt records an attachment that exists already, for example the one a worker +// started on. +func (f *Follower) Adopt(c cluster.Carrier, a Attachment, epoch int64) { + f.mu.Lock() + defer f.mu.Unlock() + f.attached[c] = a + if f.view.Epoch == 0 { + f.view.Epoch = epoch + } +} + +// Attach attaches to c with the credential that a registration returned, adopts +// the attachment, and records what the answer said about the cluster. A worker +// does this once at start, and the follower does it by itself afterwards. +func (f *Follower) Attach(ctx context.Context, c cluster.Carrier, mgr *workerregistry.CredentialManager, res *workerregistry.RegisterResponse) (Attachment, error) { + attach := f.o.Attachers[c] + if attach == nil { + return nil, fmt.Errorf("this worker has no way to use the %s carrier", c) + } + h := handoverOf(c, mgr, res) + a, err := attach(ctx, h) + if err != nil { + return nil, err + } + f.Adopt(c, a, res.CarrierEpoch) + f.Observe(beatOf(res)) + return a, nil +} + +func handoverOf(c cluster.Carrier, mgr *workerregistry.CredentialManager, res *workerregistry.RegisterResponse) Handover { + h := Handover{Carrier: c, Creds: mgr} + if c == cluster.CarrierNATS { + h.NATSURL, h.NATSCAPEM, h.NATSClientTLS = res.NatsURL, res.NatsCA, res.NatsClientTLS + } + return h +} + +// Observe records what the frontend said about the carrier, from a registration. +func (f *Follower) Observe(b *Beat) { + if b == nil || b.Carrier == "" { + return + } + f.mu.Lock() + f.view = CarrierView{ + Known: true, + Active: cluster.Carrier(b.Carrier), + Epoch: b.CarrierEpoch, + State: cluster.State(b.CarrierState), + Target: cluster.Carrier(b.CarrierTarget), + Draining: cluster.Carrier(b.CarrierDraining), + NATSClientTLS: b.NatsClientTLS, + } + f.mu.Unlock() + select { + case f.kick <- struct{}{}: + default: + } +} + +// View returns what the frontend last said. +func (f *Follower) View() CarrierView { + f.mu.Lock() + defer f.mu.Unlock() + return f.view +} + +// Attached returns the carriers that are connected now, in a fixed order. +func (f *Follower) Attached() []cluster.Carrier { + f.mu.Lock() + defer f.mu.Unlock() + return f.connectedLocked() +} + +func (f *Follower) connectedLocked() []cluster.Carrier { + var out []cluster.Carrier + for _, c := range []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel} { + if a := f.attached[c]; a != nil && a.Connected() { + out = append(out, c) + } + } + return out +} + +// Ready is nil while at least one carrier is connected. A worker whose carriers +// are all down still sends its heartbeat to the frontend, and the registry would +// show a node that looks healthy and cannot be reached. +func (f *Follower) Ready() error { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.connectedLocked()) > 0 { + return nil + } + var held []string + for c := range f.attached { + held = append(held, string(c)) + } + slices.Sort(held) + return fmt.Errorf("worker carrier is down (%s): the frontend cannot reach this worker", strings.Join(held, ", ")) +} + +// Report is the carrier part of a heartbeat. +func (f *Follower) Report() map[string]any { + f.mu.Lock() + defer f.mu.Unlock() + return f.reportLocked() +} + +func (f *Follower) reportLocked() map[string]any { + attached := f.connectedLocked() + follow := slices.Clone(attached) + var reasons []string + for _, c := range []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel} { + if _, ok := f.o.Attachers[c]; !ok { + continue + } + reason := "" + if cannot := f.o.Cannot[c]; cannot != nil { + reason = cannot(f.view) + } + if reason == "" { + reason = f.errs[c] + } + held := f.attached[c] != nil + switch { + case reason == "": + if !slices.Contains(follow, c) { + follow = append(follow, c) + } + case !held: + reasons = append(reasons, fmt.Sprintf("%s: %s", c, reason)) + } + } + // A worker holds a carrier it can follow, even before it is connected, so a + // report always names something. + for c := range f.attached { + if !slices.Contains(follow, c) { + follow = append(follow, c) + } + } + slices.Sort(follow) + names := func(cs []cluster.Carrier) []string { + out := make([]string, 0, len(cs)) + for _, c := range cs { + out = append(out, string(c)) + } + return out + } + return map[string]any{ + "attached": names(attached), + "attached_epoch": f.view.Epoch, + "follow_capabilities": names(follow), + "follow_error": strings.Join(reasons, "; "), + } +} + +// Run sends the heartbeats and follows the carrier until ctx ends. It returns +// when both loops have ended. The attachments are not closed: whoever attached +// them closes them (Close). +func (f *Follower) Run(ctx context.Context) { + f.mu.Lock() + f.root = ctx + f.mu.Unlock() + var wg sync.WaitGroup + wg.Add(2) + go func() { defer wg.Done(); f.heartbeats(ctx) }() + go func() { defer wg.Done(); f.reconcileLoop(ctx) }() + wg.Wait() +} + +// Close closes every attachment. +func (f *Follower) Close() { + f.mu.Lock() + held := f.attached + f.attached = map[cluster.Carrier]Attachment{} + f.mu.Unlock() + for c, a := range held { + if err := a.Close(); err != nil { + xlog.Warn("Closing a carrier of the worker failed", "carrier", c, "error", err) + } + } +} + +func (f *Follower) heartbeats(ctx context.Context) { + wait := time.Duration(0) // the first heartbeat is at once + for { + select { + case <-ctx.Done(): + return + case <-time.After(wait): + } + f.sendBeat(ctx) + wait = f.nextBeat() + } +} + +// nextBeat is the wait to the next heartbeat. It is short while a change is under +// way or a carrier is being released. +func (f *Follower) nextBeat() time.Duration { + f.mu.Lock() + defer f.mu.Unlock() + v := f.view + busy := v.Known && (v.State != cluster.StateStable || v.Draining != "") + if !busy && v.Known { + for c := range f.attached { + if c != v.Active { + busy = true + } + } + if f.activeMissingLocked() { + busy = true + } + } + if busy { + return min(f.o.Interval, followFastBeat) + } + return f.o.Interval +} + +func (f *Follower) activeMissingLocked() bool { + return f.view.Known && f.attached[f.view.Active] == nil +} + +// sendBeat sends one heartbeat, with the report, and takes the news from the +// answer. It is skipped when no carrier is connected, as the worker has always +// done: the registry must not show a worker healthy that nobody can reach. +func (f *Follower) sendBeat(ctx context.Context) { + f.beat.Lock() + defer f.beat.Unlock() + if err := f.Ready(); err != nil { + xlog.Warn("Skipping heartbeat: " + err.Error()) + return + } + body := map[string]any{} + if f.o.Body != nil { + body = f.o.Body() + } + for k, v := range f.Report() { + body[k] = v + } + reply, err := f.o.Heartbeat(ctx, body) + if err != nil { + xlog.Warn("Heartbeat failed", "error", err) + return + } + f.Observe(reply) +} + +func (f *Follower) reconcileLoop(ctx context.Context) { + for { + next := f.reconcile(ctx) + var timer <-chan time.Time + if next > 0 { + timer = time.After(next) + } + select { + case <-ctx.Done(): + return + case <-f.kick: + case <-timer: + } + } +} + +// wanted returns the carriers the worker should hold now, active first, and the +// carriers it may keep open besides. +func wanted(v CarrierView) (hold []cluster.Carrier, keep []cluster.Carrier) { + hold = []cluster.Carrier{v.Active} + if v.Target != "" && v.State != cluster.StateStable && v.Target != v.Active { + hold = append(hold, v.Target) + } + keep = slices.Clone(hold) + if v.Draining != "" && !slices.Contains(keep, v.Draining) { + keep = append(keep, v.Draining) + } + return hold, keep +} + +// reconcile moves the attachments towards what the frontend said. It returns the +// wait until it should run again by itself, or zero. +func (f *Follower) reconcile(ctx context.Context) time.Duration { + v := f.View() + if !v.Known { + return 0 + } + hold, keep := wanted(v) + var again time.Duration + setAgain := func(d time.Duration) { + if d > 0 && (again == 0 || d < again) { + again = d + } + } + + // An attachment that ended for good is dropped, so that it is made again. + f.mu.Lock() + for c, a := range f.attached { + if fa, ok := a.(failable); ok && fa.Failed() { + xlog.Warn("A carrier of the worker ended and will be attached again", "carrier", c) + delete(f.attached, c) + go func() { _ = a.Close() }() + } + } + f.mu.Unlock() + + for _, c := range hold { + f.mu.Lock() + have := f.attached[c] != nil + due := f.retryAt[c] + f.mu.Unlock() + if have { + continue + } + if wait := time.Until(due); wait > 0 { + setAgain(wait) + continue + } + if d := f.follow(ctx, c, v); d > 0 { + setAgain(d) + } + } + + f.release(ctx, v, keep, setAgain) + return again +} + +// follow attaches to c. It returns the wait before the next attempt when it +// fails or cannot try. +func (f *Follower) follow(ctx context.Context, c cluster.Carrier, v CarrierView) time.Duration { + attach := f.o.Attachers[c] + if attach == nil { + f.setErr(c, fmt.Sprintf("this worker has no way to use the %s carrier", c)) + return followRecheck + } + if cannot := f.o.Cannot[c]; cannot != nil { + if reason := cannot(v); reason != "" { + f.setErr(c, reason) + return followRecheck + } + } + if err := f.pause(ctx, c, v); err != nil { + return 0 + } + // The cluster may have moved while the worker waited. + if now := f.View(); now.Epoch != v.Epoch { + return time.Millisecond + } + cctx, cancel := context.WithTimeout(ctx, followCredentialWait) + defer cancel() + mgr := f.o.Credentials(c) + res, err := mgr.Acquire(cctx) + if err != nil { + return f.failed(c, fmt.Errorf("getting the credential of the %s carrier: %w", c, err)) + } + a, err := attach(f.rootCtx(ctx), handoverOf(c, mgr, res)) + if err != nil { + return f.failed(c, fmt.Errorf("attaching to the %s carrier: %w", c, err)) + } + f.mu.Lock() + f.attached[c] = a + delete(f.errs, c) + f.tries[c] = 0 + f.mu.Unlock() + xlog.Info("The worker attached to a carrier", "carrier", c, "epoch", v.Epoch) + // The frontends route by what the worker reports, so tell them once the + // carrier answers, and not at the next tick. + go f.reportWhenConnected(ctx, c, a) + return 0 +} + +func (f *Follower) rootCtx(fallback context.Context) context.Context { + f.mu.Lock() + defer f.mu.Unlock() + if f.root != nil { + return f.root + } + return fallback +} + +func (f *Follower) setErr(c cluster.Carrier, reason string) { + f.mu.Lock() + f.errs[c] = reason + f.mu.Unlock() +} + +// failed records a failed attempt and returns the wait before the next one. +func (f *Follower) failed(c cluster.Carrier, err error) time.Duration { + f.mu.Lock() + defer f.mu.Unlock() + f.errs[c] = err.Error() + f.tries[c]++ + d := followBackoffMin << min(f.tries[c]-1, 5) + d = min(d, followBackoffMax) + f.retryAt[c] = time.Now().Add(d) + xlog.Warn("The worker could not follow the carrier; it keeps the one it has and tries again", "carrier", c, "retry_in", d, "error", err) + return d +} + +// pause waits a random time before the first attach to a carrier at an epoch, so +// that workers do not connect in the same instant. +func (f *Follower) pause(ctx context.Context, c cluster.Carrier, v CarrierView) error { + f.mu.Lock() + done := f.delayed[c] == v.Epoch + f.delayed[c] = v.Epoch + // The first carrier of a worker that has none is not a change to follow. + alone := len(f.attached) == 0 + f.mu.Unlock() + if done || alone || f.o.MaxDelay < 0 { + return nil + } + d := time.Duration(rand.Int64N(int64(f.o.MaxDelay) + 1)) // #nosec G404 -- a spread of reconnects, not a secret + xlog.Info("The worker follows a change of carrier", "to", c, "epoch", v.Epoch, "after", d.Round(time.Millisecond)) + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(d): + return nil + } +} + +// reportWhenConnected sends a heartbeat as soon as the new carrier answers. +func (f *Follower) reportWhenConnected(ctx context.Context, c cluster.Carrier, a Attachment) { + deadline := time.Now().Add(followCredentialWait) + for !a.Connected() { + if time.Now().After(deadline) || ctx.Err() != nil { + return + } + select { + case <-ctx.Done(): + return + case <-time.After(50 * time.Millisecond): + } + } + f.sendBeat(ctx) + select { + case f.kick <- struct{}{}: + default: + } +} + +// release closes the carriers the cluster has released: not wanted, not draining, +// and with the carrier it moved to connected. A carrier with a request on it is +// kept for a grace period, because the answer to that request goes back on it. +func (f *Follower) release(ctx context.Context, v CarrierView, keep []cluster.Carrier, setAgain func(time.Duration)) { + f.mu.Lock() + active := f.attached[v.Active] + moved := active != nil && active.Connected() + var closing []cluster.Carrier + for c, a := range f.attached { + if slices.Contains(keep, c) { + delete(f.relFrom, c) + continue + } + if !moved { + // The worker cannot follow, or has not yet. It keeps what it has. + continue + } + since, seen := f.relFrom[c] + if !seen { + since = time.Now() + f.relFrom[c] = since + } + if a.InFlight() > 0 && time.Since(since) < followDrainGrace { + setAgain(time.Second) + continue + } + closing = append(closing, c) + } + var drop []Attachment + for _, c := range closing { + drop = append(drop, f.attached[c]) + delete(f.attached, c) + delete(f.relFrom, c) + } + f.mu.Unlock() + if len(drop) == 0 { + return + } + // Tell the frontends first: a request they send after this report goes to the + // carrier that stays. + f.sendBeat(ctx) + for i, a := range drop { + if err := a.Close(); err != nil { + xlog.Warn("Closing a released carrier failed", "carrier", closing[i], "error", err) + } + xlog.Info("The worker released a carrier", "carrier", closing[i]) + } +} + +// ErrNoCredentials is returned by an attacher that was not given what it needs. +var ErrNoCredentials = errors.New("the frontend did not hand over a credential for this carrier") diff --git a/core/services/worker/follow_carrier_test.go b/core/services/worker/follow_carrier_test.go new file mode 100644 index 000000000000..9f396cc8c1be --- /dev/null +++ b/core/services/worker/follow_carrier_test.go @@ -0,0 +1,308 @@ +package worker + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/mudler/LocalAI/core/cli/workerregistry" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/workerctl" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fakeAttachment is a hold on a carrier that a spec controls. +type fakeAttachment struct { + connected atomic.Bool + inflight atomic.Int32 + failed atomic.Bool + closed atomic.Bool +} + +func (a *fakeAttachment) Connected() bool { return a.connected.Load() && !a.closed.Load() } +func (a *fakeAttachment) InFlight() int { return int(a.inflight.Load()) } +func (a *fakeAttachment) Failed() bool { return a.failed.Load() } +func (a *fakeAttachment) Close() error { a.closed.Store(true); return nil } + +// rowFrontend answers heartbeats with the row a spec sets, and records what the +// worker reported. +type rowFrontend struct { + mu sync.Mutex + beat workerregistry.HeartbeatReply + reports []map[string]any + asked []string +} + +func (f *rowFrontend) set(b workerregistry.HeartbeatReply) { + f.mu.Lock() + f.beat = b + f.mu.Unlock() +} + +func (f *rowFrontend) heartbeat(_ context.Context, body map[string]any) (*workerregistry.HeartbeatReply, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.reports = append(f.reports, body) + b := f.beat + return &b, nil +} + +func (f *rowFrontend) last() map[string]any { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.reports) == 0 { + return nil + } + return f.reports[len(f.reports)-1] +} + +func (f *rowFrontend) register(c string) workerregistry.RegisterFunc { + return func(context.Context) (*workerregistry.RegisterResponse, error) { + f.mu.Lock() + f.asked = append(f.asked, c) + f.mu.Unlock() + res := &workerregistry.RegisterResponse{ID: "node-1", Status: "healthy", Carrier: f.beat.Carrier, CredentialFor: c} + if c == "tunnel" { + res.TunnelToken = "tunnel-token" + } else { + res.NatsURL = "nats://handed:4222" + } + return res, nil + } +} + +var _ = Describe("A worker that follows the carrier of the cluster", func() { + var ( + ctx context.Context + front *rowFrontend + nats *fakeAttachment + held map[cluster.Carrier]*fakeAttachment + attaches map[cluster.Carrier]*atomic.Int32 + attachFn func(c cluster.Carrier) Attacher + cannot map[cluster.Carrier]func(CarrierView) string + f *Follower + failNext map[cluster.Carrier]*atomic.Int32 + ) + + row := func(active cluster.Carrier, epoch int64, state, target, draining string) workerregistry.HeartbeatReply { + return workerregistry.HeartbeatReply{Carrier: string(active), CarrierEpoch: epoch, CarrierState: state, CarrierTarget: target, CarrierDraining: draining} + } + attachedNow := func() []cluster.Carrier { return f.Attached() } + + BeforeEach(func() { + var cancel context.CancelFunc + ctx, cancel = context.WithCancel(context.Background()) + DeferCleanup(cancel) + front = &rowFrontend{} + front.set(row(cluster.CarrierNATS, 1, "stable", "", "")) + held = map[cluster.Carrier]*fakeAttachment{} + attaches = map[cluster.Carrier]*atomic.Int32{cluster.CarrierNATS: {}, cluster.CarrierTunnel: {}} + failNext = map[cluster.Carrier]*atomic.Int32{cluster.CarrierNATS: {}, cluster.CarrierTunnel: {}} + cannot = map[cluster.Carrier]func(CarrierView) string{} + var mu sync.Mutex + attachFn = func(c cluster.Carrier) Attacher { + return func(_ context.Context, h Handover) (Attachment, error) { + attaches[c].Add(1) + if failNext[c].Load() > 0 { + failNext[c].Add(-1) + return nil, errors.New("the server refused the connection") + } + if c == cluster.CarrierTunnel { + Expect(h.Creds.TunnelToken()).To(Equal("tunnel-token")) + } else { + Expect(h.NATSURL).To(Equal("nats://handed:4222")) + } + a := &fakeAttachment{} + a.connected.Store(true) + mu.Lock() + held[c] = a + mu.Unlock() + return a, nil + } + } + nats = &fakeAttachment{} + nats.connected.Store(true) + }) + + start := func(maxDelay time.Duration) { + f = NewFollower(FollowerOptions{ + Attachers: map[cluster.Carrier]Attacher{ + cluster.CarrierNATS: attachFn(cluster.CarrierNATS), + cluster.CarrierTunnel: attachFn(cluster.CarrierTunnel), + }, + Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { + return workerregistry.NewCredentialManagerFor(front.register(string(c)), false, string(c)) + }, + Heartbeat: front.heartbeat, + Cannot: cannot, + Interval: 100 * time.Millisecond, + MaxDelay: maxDelay, + }) + f.Adopt(cluster.CarrierNATS, nats, 1) + go f.Run(ctx) + DeferCleanup(f.Close) + } + + It("reports the carrier it holds at once, and holds no other while the cluster is stable", func() { + start(-1) + Eventually(func() map[string]any { return front.last() }, "3s").ShouldNot(BeNil()) + rep := front.last() + Expect(rep["attached"]).To(Equal([]string{"nats"})) + Expect(rep["follow_capabilities"]).To(ContainElement("nats")) + Consistently(attachedNow, "500ms").Should(Equal([]cluster.Carrier{cluster.CarrierNATS})) + Expect(attaches[cluster.CarrierTunnel].Load()).To(BeZero()) + }) + + It("attaches to the target while a change is prepared, keeps the old carrier, and reports both", func() { + start(-1) + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(attachedNow, "5s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + Eventually(func() any { return front.last()["attached"] }, "3s").Should(Equal([]string{"nats", "tunnel"})) + Expect(front.last()["attached_epoch"]).To(BeEquivalentTo(2)) + Expect(nats.closed.Load()).To(BeFalse()) + Expect(front.asked).To(ContainElement("tunnel")) + }) + + It("never changes the active carrier by itself: it attaches to what the frontend names and nothing else", func() { + start(-1) + Consistently(func() int32 { return attaches[cluster.CarrierTunnel].Load() }, "700ms").Should(BeZero()) + }) + + It("keeps the old carrier through the drain, and releases it when the cluster does, after telling the frontend", func() { + start(-1) + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(attachedNow, "5s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + front.set(row(cluster.CarrierTunnel, 3, "commit", "nats", "nats")) + front.set(row(cluster.CarrierTunnel, 4, "stable", "", "nats")) + Consistently(func() bool { return nats.closed.Load() }, "800ms").Should(BeFalse(), "the drain keeps NATS attached") + + front.set(row(cluster.CarrierTunnel, 5, "stable", "", "")) + Eventually(func() bool { return nats.closed.Load() }, "5s").Should(BeTrue()) + Eventually(attachedNow, "3s").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + Eventually(func() any { return front.last()["attached"] }, "3s").Should(Equal([]string{"tunnel"})) + }) + + It("does not close a carrier while a request runs on it, and closes it when the request ends", func() { + start(-1) + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(attachedNow, "5s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + nats.inflight.Store(1) + front.set(row(cluster.CarrierTunnel, 3, "stable", "", "")) + Consistently(func() bool { return nats.closed.Load() }, "1500ms").Should(BeFalse()) + nats.inflight.Store(0) + Eventually(func() bool { return nats.closed.Load() }, "5s").Should(BeTrue()) + }) + + It("drops the target when the change is aborted", func() { + start(-1) + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(attachedNow, "5s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + front.set(row(cluster.CarrierNATS, 3, "stable", "", "")) + Eventually(func() bool { return held[cluster.CarrierTunnel].closed.Load() }, "5s").Should(BeTrue()) + Expect(nats.closed.Load()).To(BeFalse()) + }) + + It("cannot follow: stays on its carrier after the drain, says why, and does not try", func() { + cannot[cluster.CarrierTunnel] = func(CarrierView) string { return "no frontend URL that a tunnel can use" } + start(-1) + front.set(row(cluster.CarrierTunnel, 3, "stable", "", "")) + Eventually(func() any { return front.last()["follow_error"] }, "3s").Should(ContainSubstring("no frontend URL that a tunnel can use")) + front.set(row(cluster.CarrierTunnel, 4, "stable", "", "")) + Consistently(func() bool { return nats.closed.Load() }, "800ms").Should(BeFalse(), "it keeps the only carrier it has") + Expect(attachedNow()).To(Equal([]cluster.Carrier{cluster.CarrierNATS})) + Expect(attaches[cluster.CarrierTunnel].Load()).To(BeZero()) + Expect(front.last()["follow_capabilities"]).To(Equal([]string{"nats"})) + }) + + It("reports a carrier it can use as one it can follow, and one it cannot as an error", func() { + cannot[cluster.CarrierTunnel] = func(CarrierView) string { return "" } + cannot[cluster.CarrierNATS] = func(CarrierView) string { return "needs a client certificate" } + start(-1) + Eventually(func() any { + if r := front.last(); r != nil { + return r["follow_capabilities"] + } + return nil + }, "3s").Should(Equal([]string{"nats", "tunnel"})) + }) + + It("tries again with a wait when the carrier refuses, keeps the old one meanwhile, and says why", func() { + failNext[cluster.CarrierTunnel].Store(1) + start(-1) + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(func() any { return front.last()["follow_error"] }, "5s").Should(ContainSubstring("refused")) + Expect(attachedNow()).To(Equal([]cluster.Carrier{cluster.CarrierNATS})) + Eventually(attachedNow, "10s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + Eventually(func() any { return front.last()["follow_error"] }, "3s").Should(BeEmpty()) + Expect(attaches[cluster.CarrierTunnel].Load()).To(BeEquivalentTo(2)) + }) + + It("leaves a worker alone when the frontend predates carriers", func() { + front.set(workerregistry.HeartbeatReply{}) + start(-1) + Consistently(func() int32 { return attaches[cluster.CarrierTunnel].Load() + attaches[cluster.CarrierNATS].Load() }, "700ms").Should(BeZero()) + Expect(nats.closed.Load()).To(BeFalse()) + }) + + It("waits a random time within the bound before it attaches", func() { + start(1500 * time.Millisecond) + began := time.Now() + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(attachedNow, "6s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + Expect(time.Since(began)).To(BeNumerically("<", 5*time.Second)) + }) + + It("attaches again to a carrier that ended for good", func() { + start(-1) + front.set(row(cluster.CarrierNATS, 2, "prepare", "tunnel", "")) + Eventually(attachedNow, "5s").Should(ConsistOf(cluster.CarrierNATS, cluster.CarrierTunnel)) + first := held[cluster.CarrierTunnel] + first.failed.Store(true) + Eventually(func() bool { return first.closed.Load() }, "5s").Should(BeTrue()) + Eventually(func() int32 { return attaches[cluster.CarrierTunnel].Load() }, "10s").Should(BeEquivalentTo(2)) + }) + + It("reports no carrier as ready when all are down", func() { + start(-1) + nats.connected.Store(false) + Expect(f.Ready()).To(HaveOccurred()) + nats.connected.Store(true) + Expect(f.Ready()).To(Succeed()) + }) +}) + +var _ = Describe("The requests that run on a control server", func() { + It("counts a request of the HTTP server while it runs, and not after", func() { + srv := newHTTPControlServer() + release := make(chan struct{}) + entered := make(chan struct{}) + Expect(srv.handle(verbBackendList, func(context.Context, []byte) (any, error) { + close(entered) + <-release + return map[string]string{}, nil + })).To(Succeed()) + ts := httptest.NewServer(srv) + DeferCleanup(ts.Close) + + done := make(chan struct{}) + go func() { + defer close(done) + resp, err := http.Post(ts.URL+workerctl.PathOf(workerctl.VerbBackendList), "application/json", strings.NewReader("{}")) + if err == nil { + _ = resp.Body.Close() + } + }() + Eventually(entered, "5s").Should(BeClosed()) + Expect(srv.InFlight()).To(Equal(1)) + close(release) + Eventually(done, "5s").Should(BeClosed()) + Eventually(srv.InFlight, "5s").Should(BeZero()) + }) +}) diff --git a/core/services/worker/nats_connect.go b/core/services/worker/nats_connect.go index a3db89de65af..25485701d2ec 100644 --- a/core/services/worker/nats_connect.go +++ b/core/services/worker/nats_connect.go @@ -1,12 +1,9 @@ package worker import ( - "context" "fmt" - "github.com/mudler/LocalAI/core/cli/workerregistry" "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/xlog" ) // connectNATS opens a NATS client using JWT+seed from env or registration (env wins). @@ -34,34 +31,3 @@ func connectNATS(url, envJWT, envSeed, registerJWT, registerSeed string, require } return messaging.New(url, opts...) } - -// connectNatsFor connects a worker to NATS in the way it always did. -// -// Static credentials from the environment cannot be minted again, so they are -// used as they are. Otherwise the connection takes its credentials from the -// credential manager on every connect, and a goroutine refreshes them before the -// minted JWT expires. If the refresh fails for good, the worker stops, so that it -// restarts and acquires credentials again and does not drift towards a JWT that -// it can no longer renew. -func connectNatsFor(ctx context.Context, stop context.CancelFunc, cfg *Config, credMgr *workerregistry.CredentialManager, static bool, tls messaging.TLSFiles) (*messaging.Client, error) { - if static { - return connectNATS(cfg.NatsURL, cfg.NatsJWT, cfg.NatsUserSeed, "", "", cfg.NatsAuthRequired(), tls) - } - var opts []messaging.Option - if credMgr.HasCredentials() { - opts = append(opts, messaging.WithUserJWTProvider(credMgr.Provider())) - } - if tls.Enabled() { - opts = append(opts, messaging.WithTLS(tls)) - } - client, err := messaging.New(cfg.NatsURL, opts...) - if err == nil && credMgr.HasCredentials() { - go func() { - if err := credMgr.RefreshLoop(ctx); err != nil { - xlog.Error("NATS credential refresh permanently failed; shutting down worker", "error", err) - stop() - } - }() - } - return client, err -} diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 2f624ccd20f4..3d7c86fcc07a 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -122,11 +122,6 @@ type backendSupervisor struct { // loopback. bindHost string - // follower attaches the worker to a tunnel while it runs. It is nil for a - // worker that cannot attach, and nothing calls it until the worker follows - // the carrier of the cluster. - follower *tunnelFollower - cfg *Config ml *model.ModelLoader systemState *system.SystemState diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index e84cc3ad7b75..860b24ee33e5 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" "net" "net/http" "os" @@ -20,11 +21,11 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/LocalAI/pkg/model" - "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/xlog" ) @@ -88,62 +89,56 @@ func Run(ctx *cliContext.Context, cfg *Config) error { shutdownCtx, shutdownCancel := context.WithCancel(context.Background()) defer shutdownCancel() - registrationBody := cfg.registrationBody() - natsTLS := messaging.TLSFiles{CA: cfg.NatsTLSCA, Cert: cfg.NatsTLSCert, Key: cfg.NatsTLSKey} + // The body of a registration. A worker that follows a change registers again + // and names the carrier it wants the credential of. + baseBody := cfg.registrationBody() + var routable atomic.Bool + routable.Store(cfg.hasRoutableAddress()) + registrationBody := func() map[string]any { + b := make(map[string]any, len(baseBody)+2) + maps.Copy(b, baseBody) + b["routable"] = routable.Load() + return b + } + natsLocal := NATSLocal{ + URL: cfg.NatsURL, JWT: cfg.NatsJWT, Seed: cfg.NatsUserSeed, Required: cfg.NatsAuthRequired(), + TLS: messaging.TLSFiles{CA: cfg.NatsTLSCA, Cert: cfg.NatsTLSCert, Key: cfg.NatsTLSKey}, + } // Register, and learn from the answer which carrier the cluster runs on. // // Static NATS credentials from the environment cannot be minted again, so - // the worker registers once and uses them as they are. Otherwise the - // credential manager registers, waits through admin approval, and for NATS - // it refreshes the credentials before the minted JWT expires, so that the - // connection survives the expiry through a reconnect that the client does by - // itself. For the tunnel it holds the token of the node. - var ( - res *workerregistry.RegisterResponse - credMgr = workerregistry.NewCredentialManager( - func(ctx context.Context) (*workerregistry.RegisterResponse, error) { - return regClient.RegisterFull(ctx, registrationBody) - }, - cfg.NatsAuthRequired(), - ) - staticNATS = cfg.NatsJWT != "" || cfg.NatsUserSeed != "" - regErr error + // the worker does not wait for a minted one. Otherwise the credential manager + // registers, waits through admin approval, and for NATS it refreshes the + // credentials before the minted JWT expires, so that the connection survives + // the expiry through a reconnect that the client does by itself. For the + // tunnel it holds the token of the node. + staticNATS := cfg.NatsJWT != "" || cfg.NatsUserSeed != "" + credMgr := workerregistry.NewCredentialManager( + func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + return regClient.RegisterFull(ctx, registrationBody()) + }, + cfg.NatsAuthRequired() && !staticNATS, ) - if staticNATS { - res, regErr = regClient.RegisterFullWithRetry(shutdownCtx, registrationBody, 10) - } else { - res, regErr = credMgr.Acquire(shutdownCtx) - } + res, regErr := credMgr.Acquire(shutdownCtx) if regErr != nil { return fmt.Errorf("failed to register with frontend: %w", regErr) } nodeID := res.ID - onTunnel, err := carrierOf(res.Carrier, cfg) + onTunnel, err := carrierOf(res.Carrier, res.NatsURL, cfg) if err != nil { return err } - // The token comes from the credential manager, which is updated by every - // registration. With static NATS credentials there is no manager, and a - // frontend that answers with the tunnel gives the token once. - tunnelToken := credMgr.TunnelToken - reauthorize := credMgr.Reregister - if staticNATS { - var staticToken atomic.Pointer[string] - staticToken.Store(&res.TunnelToken) - tunnelToken = func() string { return *staticToken.Load() } - reauthorize = func(ctx context.Context) error { - again, err := regClient.RegisterFull(ctx, registrationBody) - if err != nil { - return err - } - if again.TunnelToken == "" { - return fmt.Errorf("node %s registered but the tunnel credential was not minted", again.ID) - } - staticToken.Store(&again.TunnelToken) - return nil - } + bootCarrier := cluster.CarrierNATS + if onTunnel { + bootCarrier = cluster.CarrierTunnel + } else if res.NatsClientTLS && (cfg.NatsTLSCert == "" || cfg.NatsTLSKey == "") { + return errors.New("the NATS server asks for a client certificate, and this worker has none: set LOCALAI_NATS_TLS_CERT and LOCALAI_NATS_TLS_KEY") } + // Backends of a worker that booted on NATS listen on every interface, and the + // frontends dial the address it registered. Such a worker can be dialled + // whether or not the operator gave it an address. + routable.Store(cfg.hasRoutableAddress() || !onTunnel) xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cfg.RegisterTo, "carrier", cmp.Or(res.Carrier, "nats")) heartbeatInterval, err := time.ParseDuration(cfg.HeartbeatInterval) @@ -177,36 +172,31 @@ func Run(ctx *cliContext.Context, cfg *Config) error { } xlog.Info("Ephemeral staging capacity configured", "roots", ephemeralRoots, "byteLimit", byteLimit, "minFreeBytes", minFreeBytes) StartEphemeralRootsCleanup(shutdownCtx, ephemeralRoots, ephemeralCapacity, 0, 0) - // The readiness gate is created here but only armed once NATS is up and the - // backend supervisor exists, below, because the gate probes both. + // The readiness gate is created here but only armed once a carrier is up and + // the backend supervisor exists, below, because the gate probes both. // Until then /readyz reports ready, which is correct: reaching this line // means the worker has already registered with the frontend, so it is // mid-startup rather than broken. readiness := &nodes.WorkerReadiness{} - // A worker on the tunnel serves its control verbs on this server, because the - // tunnel reaches it with the http stream tag. The routes are added once the - // supervisor exists. - var ( - httpControl *httpControlServer - controlMux http.Handler - ) - // A worker on NATS gets an empty slot instead, when it could attach to a - // tunnel later: the plane is mounted at any time, and until then the path - // answers as it always did on NATS. - var controlSlot *nodes.ControlSlot + // The control verbs are served over HTTP on this server when the worker holds + // a tunnel, because the tunnel reaches it with the http stream tag. The + // control plane is a slot on the server, empty until the worker attaches to a + // tunnel and empty again when it releases it. While it is empty the path + // answers as a path that does not exist, which is what a worker on NATS has + // always answered. A worker that has no frontend URL a tunnel can use mounts + // no plane at all. follower := newTunnelFollower(tunnelFollowerOptions{ - Config: cfg, - NodeID: nodeID, - HTTPAddr: httpAddr, - Token: tunnelToken, - Reauthorize: reauthorize, + Config: cfg, + NodeID: nodeID, + HTTPAddr: httpAddr, }) - switch { - case onTunnel: - httpControl = newHTTPControlServer() - controlMux = httpControl - case follower.CanAttach(): - controlSlot = &nodes.ControlSlot{} + if onTunnel && !follower.CanAttach() { + _, terr := tunnelEndpoint(cfg.RegisterTo, nodeID) + return terr + } + var controlMux http.Handler + if follower.CanAttach() { + controlSlot := &nodes.ControlSlot{} follower.slot = controlSlot controlMux = controlSlot } @@ -215,61 +205,6 @@ func Run(ctx *cliContext.Context, cfg *Config) error { return fmt.Errorf("starting HTTP file transfer server: %w", err) } - // Attach to the carrier that the frontend named. The link is what the - // heartbeat and the readiness probe look at: a worker whose link is down - // still sends its HTTP heartbeat, and the registry would show a node that - // looks healthy and cannot be reached. - var ( - link carrierLink - natsClient *messaging.Client - linkProbe func() error - linkDownMsg string - ) - if onTunnel { - t, terr := StartTunnel(shutdownCtx, tunnelConfig(cfg, nodeID, httpAddr, tunnelToken, reauthorize)) - if terr != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("starting the worker tunnel: %w", terr) - } - defer func() { - if err := t.Close(); err != nil { - xlog.Warn("Closing the worker tunnel failed", "error", err) - } - }() - link, linkProbe, linkDownMsg = t, nodes.TunnelReadiness(t), "tunnel disconnected" - } else { - xlog.Info("Connecting to NATS", "url", sanitize.URL(cfg.NatsURL)) - var cerr error - natsClient, cerr = connectNatsFor(shutdownCtx, shutdownCancel, cfg, credMgr, staticNATS, natsTLS) - if cerr != nil { - nodes.ShutdownFileTransferServer(httpServer) - return fmt.Errorf("connecting to NATS: %w", cerr) - } - defer natsClient.Close() - link, linkProbe, linkDownMsg = natsLink{natsClient}, nodes.NATSReadiness(natsClient), "NATS disconnected" - } - - // Start heartbeat goroutine (after the carrier is attached so the check works) - go func() { - ticker := time.NewTicker(heartbeatInterval) - defer ticker.Stop() - for { - select { - case <-shutdownCtx.Done(): - return - case <-ticker.C: - if !link.Connected() { - xlog.Warn("Skipping heartbeat: " + linkDownMsg) - continue - } - body := cfg.heartbeatBody() - if err := regClient.Heartbeat(shutdownCtx, nodeID, body); err != nil { - xlog.Warn("Heartbeat failed", "error", err) - } - } - } - }() - // Process supervisor — manages multiple backend gRPC processes on different ports basePort := cfg.effectiveBasePort() // Buffered so NATS stop handler can send without blocking @@ -299,15 +234,50 @@ func Run(ctx *cliContext.Context, cfg *Config) error { bindHost: bindHost, } + // registerVerbs claims every verb of the worker on a control server: the + // NATS one, or the HTTP one behind the tunnel. It is called for each carrier + // the worker attaches to, at start and when the cluster changes. + var registerVerbs func(control controlServer) error + + // The carriers of the worker. It follows the one the frontend names, so the + // attachments are made by the same code at start and when the cluster changes. + carriers := NewFollower(FollowerOptions{ + Attachers: map[cluster.Carrier]Attacher{ + cluster.CarrierNATS: natsAttacher(natsLocal, nodeID, func(cs controlServer) error { return registerVerbs(cs) }), + cluster.CarrierTunnel: tunnelAttacher(follower), + }, + Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { + return workerregistry.NewCredentialManagerFor( + registerFor(regClient, registrationBody, c), c == cluster.CarrierNATS && cfg.NatsAuthRequired() && !staticNATS, string(c)) + }, + Heartbeat: func(ctx context.Context, body map[string]any) (*Beat, error) { + return regClient.HeartbeatFull(ctx, nodeID, body) + }, + Body: cfg.heartbeatBody, + Cannot: map[cluster.Carrier]func(CarrierView) string{ + cluster.CarrierNATS: NATSReasons(routable.Load(), natsLocal), + cluster.CarrierTunnel: func(CarrierView) string { + if !follower.CanAttach() { + return "this worker has no frontend URL that a tunnel can use" + } + return "" + }, + }, + Interval: heartbeatInterval, + MaxDelay: cfg.followMaxDelay(), + }) + defer carriers.Close() + // Arm the readiness gate now that the worker can actually receive work. - // The carrier is already attached at this point, so a worker that is up but - // cut off from it reports 503 instead of a meaningless 200 (#10987). + // The carrier is attached before the worker serves, so a worker that is up + // but cut off from every carrier reports 503 instead of a meaningless 200 + // (#10987). // // Readiness also covers the data path: a worker whose link is fine but whose // backend processes have died is up and useless, and the scheduler cannot // tell the difference from the link alone. readiness.Set(nodes.CompositeReadiness( - linkProbe, + carriers.Ready, nodes.BackendDataPathReadiness(supervisor, func(addr string) error { conn, err := net.DialTimeout("tcp", addr, 2*time.Second) if err != nil { @@ -327,10 +297,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // The watchdog stops load operations the controller no longer renews. go supervisor.runOperationWatchdog(shutdownCtx) - // registerVerbs claims every verb of the worker on a control server: the - // NATS one, the HTTP one of a worker on the tunnel, or the HTTP one that a - // worker on NATS gains when it attaches to a tunnel. - registerVerbs := func(control controlServer) error { + registerVerbs = func(control controlServer) error { if err := supervisor.registerLifecycleVerbs(control); err != nil { return fmt.Errorf("serving the worker lifecycle verbs: %w", err) } @@ -342,58 +309,27 @@ func Run(ctx *cliContext.Context, cfg *Config) error { } return nil } - var control controlServer - if natsClient != nil { - control = newNATSControlServer(natsClient, nodeID) - } else { - control = httpControl - } - if err := registerVerbs(control); err != nil { + follower.register = registerVerbs + + // Attach to the carrier that the frontend named, with the credential that the + // registration returned, and follow it from here on. + if _, err := carriers.Attach(shutdownCtx, bootCarrier, credMgr, res); err != nil { nodes.ShutdownFileTransferServer(httpServer) - return err - } - // The hook of a worker on NATS that attaches to a tunnel later. Nothing calls - // it yet, so a worker on NATS behaves as it did. - if controlSlot != nil { - follower.register = registerVerbs - supervisor.follower = follower - defer func() { - if err := follower.Detach(); err != nil { - xlog.Warn("Closing the worker tunnel failed", "error", err) - } - }() + return fmt.Errorf("attaching to the %s carrier: %w", bootCarrier, err) } + go carriers.Run(shutdownCtx) xlog.Info("Worker ready, waiting for backend.install events") - // Exit on an OS signal or on an internal fatal condition (e.g. NATS - // credentials became unrenewable), so the worker restarts and re-acquires - // rather than lingering unable to serve. - var runErr error - select { - case <-sigCh: - case <-shutdownCtx.Done(): - runErr = fmt.Errorf("worker shutting down: NATS credentials unavailable") - xlog.Error("Internal shutdown requested", "error", runErr) - } + <-sigCh xlog.Info("Shutting down worker") shutdownCancel() // stop heartbeat loop immediately regClient.GracefulDeregister(nodeID) supervisor.stopAllBackends(false) nodes.ShutdownFileTransferServer(httpServer) - return runErr -} - -// carrierLink is the connection of a worker to the carrier it is attached to. -type carrierLink interface { - Connected() bool + return nil } -// natsLink adapts a NATS client to carrierLink. -type natsLink struct{ c *messaging.Client } - -func (l natsLink) Connected() bool { return l.c.IsConnected() } - // carrierOf maps the carrier that the frontend named to the one this worker // attaches to, and reports whether it is the tunnel. // @@ -401,19 +337,19 @@ func (l natsLink) Connected() bool { return l.c.IsConnected() } // Then the worker needs a NATS URL, and without one it stops with a message that // says so. The check is here and not in the flag, because the flag cannot know // the carrier: for a cluster on the tunnel the URL is not needed. -func carrierOf(named string, cfg *Config) (onTunnel bool, err error) { +func carrierOf(named, handedNATSURL string, cfg *Config) (onTunnel bool, err error) { switch named { case "tunnel": if cfg.NatsURL != "" { - xlog.Info("The cluster runs on the tunnel, so the NATS URL of this worker is not used") + xlog.Info("The cluster runs on the tunnel; the NATS URL of this worker is kept for a change to NATS") } return true, nil case "", "nats": - if cfg.NatsURL == "" { + if cfg.NatsURL == "" && handedNATSURL == "" { if named == "" { return false, errors.New("the frontend does not name a carrier, so the cluster runs on NATS, and this worker has no NATS URL: set LOCALAI_NATS_URL") } - return false, errors.New("the cluster runs on NATS and this worker has no NATS URL: set LOCALAI_NATS_URL") + return false, errors.New("the cluster runs on NATS, this worker has no NATS URL and the frontend handed over none: set LOCALAI_NATS_URL, or store the NATS address in the cluster settings") } return false, nil default: From 82696409edc9470b9825a7cd006e87e84618dba3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 13:55:52 +0000 Subject: [PATCH 125/162] feat(agentworker): follow the carrier that the cluster names An agent worker serves its runs, its MCP requests and its jobs on the carrier the cluster uses, and now follows a change as a backend worker does. The handlers are the same on both carriers, and the worker attaches them to NATS or to the tunnel as the frontend names the carrier: on NATS it connects with the credential and the address the frontend hands over (a local address or credential still wins), consumes the queues, and serves the MCP subjects; on the tunnel it serves the same verbs behind its loopback server. A run keeps the carrier it began on, and the carrier is not closed while a run is in flight on it. The agent worker opens no port, so the only thing that can stop it from following is a NATS server that asks for a client certificate, which cannot be handed over. It reports that, and its heartbeat carries the carriers it is attached to like the heartbeat of a backend worker. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/agent_worker.go | 286 ++++++++------------------ core/cli/agent_worker_subject_test.go | 16 +- core/services/agentworker/follow.go | 227 ++++++++++++++++++++ core/services/worker/config.go | 2 +- core/services/worker/follow_attach.go | 9 + 5 files changed, 333 insertions(+), 207 deletions(-) create mode 100644 core/services/agentworker/follow.go diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 8f62a8c790ef..b0c58e7dd807 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -15,13 +15,12 @@ import ( "github.com/mudler/LocalAI/core/cli/workerregistry" "github.com/mudler/LocalAI/core/config" mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" - "github.com/mudler/LocalAI/core/services/agents" "github.com/mudler/LocalAI/core/services/agentworker" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/jobs" mcpRemote "github.com/mudler/LocalAI/core/services/mcp" "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/LocalAI/core/services/nodes" - "github.com/mudler/LocalAI/core/services/workerctl" + "github.com/mudler/LocalAI/core/services/worker" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/cogito" @@ -41,8 +40,8 @@ import ( // // localai agent-worker --nats-url nats://... --register-to http://localai:8080 type AgentWorkerCMD struct { - // NATS (required) - NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed only while the cluster runs on NATS" group:"distributed"` + // NATS (optional: the frontend hands over the address) + NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Optional: the frontend hands over the address when the cluster runs on NATS, and this overrides it" group:"distributed"` // Registration (required) RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"` @@ -97,16 +96,16 @@ func validateAgentSubject(subject string) error { // none predates carriers and runs on NATS. The flag of the NATS URL cannot be // required, because the carrier is known only after the registration, so the // check is here. -func agentCarrier(named, natsURL string) (onTunnel bool, err error) { +func agentCarrier(named, natsURL, handedURL string) (onTunnel bool, err error) { switch named { case "tunnel": if natsURL != "" { - xlog.Info("The cluster runs on the tunnel, so the NATS URL of this agent worker is not used") + xlog.Info("The cluster runs on the tunnel; the NATS URL of this agent worker is kept for a change to NATS") } return true, nil case "", "nats": - if natsURL == "" { - return false, fmt.Errorf("the cluster runs on NATS and this agent worker has no NATS URL: set LOCALAI_NATS_URL") + if natsURL == "" && handedURL == "" { + return false, fmt.Errorf("the cluster runs on NATS and this agent worker has no NATS URL, and the frontend handed over none: set LOCALAI_NATS_URL") } return false, nil default: @@ -134,14 +133,17 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { hostname, _ := os.Hostname() nodeName = "agent-" + hostname } - registrationBody := map[string]any{ - "name": nodeName, - "node_type": "agent", - "version": internal.Version, - "commit": internal.Commit, - } - if cmd.RegistrationToken != "" { - registrationBody["token"] = cmd.RegistrationToken + registrationBody := func() map[string]any { + body := map[string]any{ + "name": nodeName, + "node_type": "agent", + "version": internal.Version, + "commit": internal.Commit, + } + if cmd.RegistrationToken != "" { + body["token"] = cmd.RegistrationToken + } + return body } // Context cancelled on shutdown — used by registration waits, heartbeat, and @@ -152,11 +154,12 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // Acquire credentials via (re)registration. When the bus requires auth and no // static fallback is configured, wait through admin approval until the // frontend mints credentials rather than starting unauthenticated. + staticNATS := cmd.NatsJWT != "" || cmd.NatsServiceJWT != "" credMgr := workerregistry.NewCredentialManager( func(ctx context.Context) (*workerregistry.RegisterResponse, error) { - return regClient.RegisterFull(ctx, registrationBody) + return regClient.RegisterFull(ctx, registrationBody()) }, - cmd.natsAuthRequired() && cmd.NatsJWT == "" && cmd.NatsServiceJWT == "", + cmd.natsAuthRequired() && !staticNATS, ) res, err := credMgr.Acquire(shutdownCtx) if err != nil { @@ -165,141 +168,96 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { nodeID := res.ID xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo) - onTunnel, err := agentCarrier(res.Carrier, cmd.NatsURL) + onTunnel, err := agentCarrier(res.Carrier, cmd.NatsURL, res.NatsURL) if err != nil { return err } + bootCarrier := cluster.CarrierNATS + if onTunnel { + bootCarrier = cluster.CarrierTunnel + } // Use provisioned API token if none was set if cmd.APIToken == "" { cmd.APIToken = res.APIToken } - // Start heartbeat heartbeatInterval, err := time.ParseDuration(cmd.HeartbeatInterval) if err != nil && cmd.HeartbeatInterval != "" { xlog.Warn("invalid heartbeat interval, using default 10s", "input", cmd.HeartbeatInterval, "error", err) } heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second) - go regClient.HeartbeatLoop(shutdownCtx, nodeID, heartbeatInterval, func() map[string]any { return map[string]any{} }) - - if onTunnel { - return cmd.runOnTunnel(shutdownCtx, shutdownCancel, regClient, credMgr, nodeID, apiURL) - } - - // Resolve NATS credentials with precedence: explicit env override, then - // frontend-minted (auto-refreshed before expiry), then service fallback. - // Each static source must supply JWT and seed together. - natsTLS := messaging.TLSFiles{CA: cmd.NatsTLSCA, Cert: cmd.NatsTLSCert, Key: cmd.NatsTLSKey} - var natsOpts []messaging.Option - switch { - case cmd.NatsJWT != "" || cmd.NatsUserSeed != "": - if (cmd.NatsJWT == "") != (cmd.NatsUserSeed == "") { - return fmt.Errorf("LOCALAI_NATS_JWT and LOCALAI_NATS_USER_SEED must be set together") - } - natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsJWT, cmd.NatsUserSeed)) - case credMgr.HasCredentials(): - natsOpts = append(natsOpts, messaging.WithUserJWTProvider(credMgr.Provider())) - go func() { - if err := credMgr.RefreshLoop(shutdownCtx); err != nil { - xlog.Error("NATS credential refresh permanently failed; shutting down agent worker", "error", err) - shutdownCancel() - } - }() - case cmd.NatsServiceJWT != "" || cmd.NatsServiceSeed != "": - if (cmd.NatsServiceJWT == "") != (cmd.NatsServiceSeed == "") { - return fmt.Errorf("LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED must be set together") - } - natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsServiceJWT, cmd.NatsServiceSeed)) - case cmd.natsAuthRequired(): - return fmt.Errorf("NATS JWT+seed required: enable frontend minting or set LOCALAI_NATS_* env vars") - } - if natsTLS.Enabled() { - natsOpts = append(natsOpts, messaging.WithTLS(natsTLS)) - } - natsClient, err := messaging.New(cmd.NatsURL, natsOpts...) - if err != nil { - return fmt.Errorf("connecting to NATS: %w", err) - } - defer natsClient.Close() - - // Create event bridge for publishing results back via NATS - eventBridge := agents.NewEventBridge(natsClient, nil, "agent-worker-"+nodeID) - - // Start cancel listener - cancelSub, err := eventBridge.StartCancelListener() - if err != nil { - xlog.Warn("Failed to start cancel listener", "error", err) - } else { - defer cancelSub.Unsubscribe() - } - - // One consumer serves both queued kinds; the route option only moves the - // agent-run subject and group, which operators may set. - work := messaging.NewNATSWorkConsumer(natsClient, messaging.WithAgentRunRoute(cmd.Subject, cmd.Queue)) - - // Create and start the NATS dispatcher. - // No ConfigProvider or SkillStore needed — config and skills arrive in the job payload. - dispatcher := agents.NewNATSDispatcher( - work, - eventBridge, - nil, // no ConfigProvider: config comes in the enriched NATS payload - apiURL, cmd.APIToken, - 0, // no concurrency limit (CLI worker) - ) - - if err := dispatcher.Start(shutdownCtx); err != nil { - return fmt.Errorf("starting dispatcher: %w", err) - } - - var rpc agentRPCServer = nodes.NewNATSAgentRPCServer(natsClient, nodeID) - - // Serve MCP tool execution requests (load-balanced across workers). - // The frontend routes model-level MCP tool calls here. - if err := rpc.ServeMCPTool(handleMCPToolRequest); err != nil { - return err - } - - // Serve MCP discovery requests (load-balanced across workers). - if err := rpc.ServeMCPDiscovery(handleMCPDiscoveryRequest); err != nil { - return err - } - - // Subscribe to MCP CI job execution (load-balanced across agent workers). - // In distributed mode, MCP CI jobs are routed here because the frontend - // cannot create MCP sessions (e.g., stdio servers using docker). mcpCIJobTimeout, err := time.ParseDuration(cmd.MCPCIJobTimeout) if err != nil && cmd.MCPCIJobTimeout != "" { xlog.Warn("invalid MCP CI job timeout, using default 10m", "input", cmd.MCPCIJobTimeout, "error", err) } mcpCIJobTimeout = cmp.Or(mcpCIJobTimeout, config.DefaultMCPCIJobTimeout) - if _, err := startMCPCIConsumer(shutdownCtx, work, apiURL, cmd.APIToken, mcpCIJobTimeout); err != nil { - return err + // Resolve NATS credentials with precedence: explicit env override, then + // frontend-minted (auto-refreshed before expiry), then service fallback. + // Each static source must supply JWT and seed together. + if (cmd.NatsJWT == "") != (cmd.NatsUserSeed == "") { + return fmt.Errorf("LOCALAI_NATS_JWT and LOCALAI_NATS_USER_SEED must be set together") + } + if (cmd.NatsServiceJWT == "") != (cmd.NatsServiceSeed == "") { + return fmt.Errorf("LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED must be set together") + } + local := worker.NATSLocal{ + URL: cmd.NatsURL, JWT: cmd.NatsJWT, Seed: cmd.NatsUserSeed, Required: cmd.natsAuthRequired(), + TLS: messaging.TLSFiles{CA: cmd.NatsTLSCA, Cert: cmd.NatsTLSCert, Key: cmd.NatsTLSKey}, + FallbackJWT: cmd.NatsServiceJWT, + FallbackSeed: cmd.NatsServiceSeed, + } + serve := agentworker.FollowConfig{ + NodeID: nodeID, FrontendURL: cmd.RegisterTo, ControlToken: cmd.RegistrationToken, + Local: local, Subject: cmd.Subject, Queue: cmd.Queue, + APIURL: apiURL, APIToken: cmd.APIToken, + // Model-level MCP tool calls and discovery from the frontend. The worker + // creates and caches the MCP sessions from the serialized config. + Tool: handleMCPToolRequest, + Discovery: handleMCPDiscoveryRequest, + // The agent worker has no model loader, so it drops the sessions cached for + // a backend that goes away when it hears the stop of that backend. + BackendStop: func(backend string) { + if backend != "" { + mcpTools.CloseMCPSessions(backend) + } + }, + // MCP CI jobs run here because the frontend cannot create MCP sessions + // (for example stdio servers that use docker). + Jobs: func(ctx context.Context, consumer messaging.WorkConsumer) error { + _, err := startMCPCIConsumer(ctx, consumer, apiURL, cmd.APIToken, mcpCIJobTimeout) + return err + }, } - // Listen for backend stop events to clean up cached MCP sessions. - // In the main application this is done via ml.OnModelUnload, but the agent - // worker has no model loader, so it listens for the stop event instead. - if err := rpc.ServeBackendStop(func(backend string) { - if backend != "" { - mcpTools.CloseMCPSessions(backend) - } - }); err != nil { - return err + carriers := worker.NewFollower(worker.FollowerOptions{ + Attachers: serve.Attachers(), + Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { + return workerregistry.NewCredentialManagerFor( + func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + body := registrationBody() + body["carrier"] = string(c) + return regClient.RegisterFull(ctx, body) + }, c == cluster.CarrierNATS && cmd.natsAuthRequired() && !staticNATS, string(c)) + }, + Heartbeat: func(ctx context.Context, body map[string]any) (*workerregistry.HeartbeatReply, error) { + return regClient.HeartbeatFull(ctx, nodeID, body) + }, + Cannot: serve.Cannot(), + Interval: heartbeatInterval, + }) + defer carriers.Close() + if _, err := carriers.Attach(shutdownCtx, bootCarrier, credMgr, res); err != nil { + return fmt.Errorf("attaching to the %s carrier: %w", bootCarrier, err) } + go carriers.Run(shutdownCtx) + xlog.Info("Agent worker ready, waiting for jobs", "carrier", bootCarrier, "subject", cmd.Subject, "queue", cmd.Queue) - xlog.Info("Agent worker ready, waiting for jobs", "subject", cmd.Subject, "queue", cmd.Queue) - - // Wait for an OS signal or an internal fatal condition (e.g. NATS - // credentials became unrenewable), so the worker restarts and re-acquires - // rather than lingering unable to serve. - runErr := awaitShutdown(shutdownCtx, "NATS credentials unavailable") - + runErr := awaitShutdown(shutdownCtx, "the carrier stopped") xlog.Info("Shutting down agent worker") shutdownCancel() // stop heartbeat loop immediately - _ = dispatcher.Stop() mcpTools.CloseAllMCPSessions() regClient.GracefulDeregister(nodeID) return runErr @@ -323,72 +281,6 @@ func awaitShutdown(shutdownCtx context.Context, fatal string) error { } } -// runOnTunnel serves the agent worker on the tunnel carrier: no connection to a -// bus, one HTTP server on loopback that the frontend reaches through the tunnel -// of this worker, and the same handlers as on NATS. The frontend picks the worker -// and sends the request, so the dispatcher below gets its deliveries from a -// consumer that the control server feeds. -func (cmd *AgentWorkerCMD) runOnTunnel(shutdownCtx context.Context, shutdownCancel context.CancelFunc, regClient *workerregistry.RegistrationClient, credMgr *workerregistry.CredentialManager, nodeID, apiURL string) error { - work := agentworker.NewWork() - - // No bus, so no broadcaster: the events of a run are published on the - // publisher of its delivery, which writes them to the frontend that drives the - // run. The cancel of a run is its request: the frontend ends the stream. - eventBridge := agents.NewEventBridge(nil, nil, "agent-worker-"+nodeID) - dispatcher := agents.NewNATSDispatcher(work, eventBridge, nil, apiURL, cmd.APIToken, 0) - // The consumers are registered before the tunnel starts, so that a frontend - // which connects at once finds a worker that can take its first run. - if err := dispatcher.Start(shutdownCtx); err != nil { - return fmt.Errorf("starting dispatcher: %w", err) - } - - mcpCIJobTimeout, err := time.ParseDuration(cmd.MCPCIJobTimeout) - if err != nil && cmd.MCPCIJobTimeout != "" { - xlog.Warn("invalid MCP CI job timeout, using default 10m", "input", cmd.MCPCIJobTimeout, "error", err) - } - mcpCIJobTimeout = cmp.Or(mcpCIJobTimeout, config.DefaultMCPCIJobTimeout) - if _, err := startMCPCIConsumer(shutdownCtx, work, apiURL, cmd.APIToken, mcpCIJobTimeout); err != nil { - return err - } - - runtime, err := agentworker.Start(shutdownCtx, agentworker.Options{ - FrontendURL: cmd.RegisterTo, - NodeID: nodeID, - // The credential is new after every registration, so it is read at every - // dial and not once. - TunnelToken: credMgr.TunnelToken, - Reauthorize: credMgr.Reregister, - ControlToken: cmd.RegistrationToken, - Handler: agentworker.Handler(agentworker.Config{ - MCPTool: agentworker.Unary(handleMCPToolRequest), - MCPDiscovery: agentworker.Unary(handleMCPDiscoveryRequest), - // The node's backend stop: the agent worker has no model loader, so it - // drops the sessions cached for the backend that goes away. - BackendStop: func(_ context.Context, req workerctl.BackendStopRequest) error { - if req.Backend != "" { - mcpTools.CloseMCPSessions(req.Backend) - } - return nil - }, - }, work), - }) - if err != nil { - return err - } - xlog.Info("Agent worker ready, waiting for runs over its tunnel") - - runErr := awaitShutdown(shutdownCtx, "the control plane stopped") - xlog.Info("Shutting down agent worker") - shutdownCancel() - dispatcher.Stop() - if err := runtime.Close(); err != nil { - xlog.Warn("Closing the agent worker tunnel failed", "error", err) - } - mcpTools.CloseAllMCPSessions() - regClient.GracefulDeregister(nodeID) - return runErr -} - // startMCPCIConsumer serves MCP CI jobs with maxInFlight 1, which keeps them // one at a time per worker, run inline on the delivery, as they always were. func startMCPCIConsumer(ctx context.Context, consumer messaging.WorkConsumer, apiURL, apiToken string, jobTimeout time.Duration) (messaging.Subscription, error) { @@ -397,14 +289,6 @@ func startMCPCIConsumer(ctx context.Context, consumer messaging.WorkConsumer, ap }) } -// agentRPCServer is how the agent worker serves the frontend's MCP requests -// and hears the node's backend stop events. -type agentRPCServer interface { - ServeMCPTool(h mcpRemote.ToolHandler) error - ServeMCPDiscovery(h mcpRemote.DiscoveryHandler) error - ServeBackendStop(h func(backend string)) error -} - // handleMCPToolRequest executes one MCP tool call. The worker creates/caches // MCP sessions from the serialized config and executes the tool. func handleMCPToolRequest(parent context.Context, req mcpRemote.MCPToolRequest) mcpRemote.MCPToolResponse { diff --git a/core/cli/agent_worker_subject_test.go b/core/cli/agent_worker_subject_test.go index 6bf175c461aa..012c8abbef94 100644 --- a/core/cli/agent_worker_subject_test.go +++ b/core/cli/agent_worker_subject_test.go @@ -33,29 +33,35 @@ var _ = Describe("validateAgentSubject", func() { var _ = Describe("agentCarrier", func() { It("attaches to NATS when the frontend names it or names nothing, and the URL is there", func() { for _, named := range []string{"nats", ""} { - onTunnel, err := agentCarrier(named, "nats://n:4222") + onTunnel, err := agentCarrier(named, "nats://n:4222", "") Expect(err).ToNot(HaveOccurred(), named) Expect(onTunnel).To(BeFalse(), named) } }) It("asks for the URL when the cluster runs on NATS and there is none", func() { - _, err := agentCarrier("nats", "") + _, err := agentCarrier("nats", "", "") Expect(err).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) - _, err = agentCarrier("", "") + _, err = agentCarrier("", "", "") Expect(err).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) }) + It("takes the URL the frontend hands over when the worker has none", func() { + onTunnel, err := agentCarrier("nats", "", "nats://handed:4222") + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(BeFalse()) + }) + It("attaches to the tunnel when the frontend names it, with a URL or without one", func() { for _, url := range []string{"", "nats://n:4222"} { - onTunnel, err := agentCarrier("tunnel", url) + onTunnel, err := agentCarrier("tunnel", url, "") Expect(err).ToNot(HaveOccurred()) Expect(onTunnel).To(BeTrue()) } }) It("asks for an upgrade when the frontend names a carrier it does not know", func() { - _, err := agentCarrier("smoke-signals", "nats://n:4222") + _, err := agentCarrier("smoke-signals", "nats://n:4222", "") Expect(err).To(MatchError(ContainSubstring("upgrade"))) }) }) diff --git a/core/services/agentworker/follow.go b/core/services/agentworker/follow.go new file mode 100644 index 000000000000..b3fb4819d13b --- /dev/null +++ b/core/services/agentworker/follow.go @@ -0,0 +1,227 @@ +package agentworker + +import ( + "context" + "fmt" + "sync/atomic" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/cluster" + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/worker" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// FollowConfig is what an agent worker needs to serve its work on either carrier. +// The handlers are the same on both: a worker follows a change of carrier by +// attaching them to the other one. +type FollowConfig struct { + NodeID string + // FrontendURL and ControlToken are for the tunnel. See Options. + FrontendURL string + ControlToken string + // Local is what the operator set for NATS. + Local worker.NATSLocal + // Subject and Queue route the runs of agents on NATS. + Subject, Queue string + // APIURL and APIToken are how a run calls back to the frontend. + APIURL, APIToken string + // Tool and Discovery answer the MCP requests of the frontend. + Tool func(context.Context, mcpRemote.MCPToolRequest) mcpRemote.MCPToolResponse + Discovery func(context.Context, mcpRemote.MCPDiscoveryRequest) mcpRemote.MCPDiscoveryResponse + // BackendStop drops the sessions cached for a backend that goes away. + BackendStop func(backend string) + // Jobs registers the consumers of the queued jobs on the consumer of one + // carrier. It is called for each attachment. + Jobs func(ctx context.Context, consumer messaging.WorkConsumer) error +} + +// counted counts the runs that are in flight on one attachment. +type counted struct { + messaging.WorkConsumer + n *atomic.Int32 +} + +func (c counted) Consume(ctx context.Context, kind messaging.WorkKind, maxInFlight int, h messaging.WorkHandler) (messaging.Subscription, error) { + return c.WorkConsumer.Consume(ctx, kind, maxInFlight, func(ctx context.Context, payload []byte, events messaging.Publisher) error { + c.n.Add(1) + defer c.n.Add(-1) + return h(ctx, payload, events) + }) +} + +// Attachers returns the attachers of an agent worker for each carrier. They are +// what the follower of the worker calls at start and when the cluster changes. +func (c FollowConfig) Attachers() map[cluster.Carrier]worker.Attacher { + return map[cluster.Carrier]worker.Attacher{ + cluster.CarrierNATS: c.natsAttacher, + cluster.CarrierTunnel: c.tunnelAttacher, + } +} + +// Cannot says why an agent worker cannot attach to a carrier. It opens no port, +// so the only thing it can lack for NATS is a certificate that cannot be handed +// over. +func (c FollowConfig) Cannot() map[cluster.Carrier]func(worker.CarrierView) string { + return map[cluster.Carrier]func(worker.CarrierView) string{ + cluster.CarrierNATS: func(v worker.CarrierView) string { + if v.NATSClientTLS && (c.Local.TLS.Cert == "" || c.Local.TLS.Key == "") { + return "the NATS server asks for a client certificate, and this worker has none: set LOCALAI_NATS_TLS_CERT and LOCALAI_NATS_TLS_KEY" + } + return "" + }, + } +} + +type agentAttachment struct { + connected func() bool + runs *atomic.Int32 + failed *atomic.Bool + close func() error +} + +func (a *agentAttachment) Connected() bool { return a.connected() } +func (a *agentAttachment) InFlight() int { return int(a.runs.Load()) } +func (a *agentAttachment) Failed() bool { return a.failed != nil && a.failed.Load() } +func (a *agentAttachment) Close() error { return a.close() } + +func (c FollowConfig) natsAttacher(ctx context.Context, h worker.Handover) (worker.Attachment, error) { + actx, cancel := context.WithCancel(ctx) + failed := &atomic.Bool{} + client, err := worker.ConnectNATS(actx, c.Local, h, func() { failed.Store(true) }) + if err != nil { + cancel() + return nil, err + } + runs := &atomic.Int32{} + var undo []func() + fail := func(err error) (worker.Attachment, error) { + for i := len(undo) - 1; i >= 0; i-- { + undo[i]() + } + client.Close() + cancel() + return nil, err + } + + // The events of a run go back over the bus the run came on. + bridge := agents.NewEventBridge(client, nil, "agent-worker-"+c.NodeID) + cancelSub, err := bridge.StartCancelListener() + if err != nil { + xlog.Warn("Failed to start cancel listener", "error", err) + } else { + undo = append(undo, func() { _ = cancelSub.Unsubscribe() }) + } + // One consumer serves both queued kinds; the route option only moves the + // agent-run subject and group, which operators may set. + consumer := counted{messaging.NewNATSWorkConsumer(client, messaging.WithAgentRunRoute(c.Subject, c.Queue)), runs} + dispatcher := agents.NewNATSDispatcher(consumer, bridge, nil, c.APIURL, c.APIToken, 0) + if err := dispatcher.Start(actx); err != nil { + return fail(fmt.Errorf("starting dispatcher: %w", err)) + } + undo = append(undo, func() { _ = dispatcher.Stop() }) + + rpc := nodes.NewNATSAgentRPCServer(client, c.NodeID) + if c.Tool != nil { + if err := rpc.ServeMCPTool(c.Tool); err != nil { + return fail(err) + } + } + if c.Discovery != nil { + if err := rpc.ServeMCPDiscovery(c.Discovery); err != nil { + return fail(err) + } + } + if c.Jobs != nil { + if err := c.Jobs(actx, consumer); err != nil { + return fail(err) + } + } + if c.BackendStop != nil { + if err := rpc.ServeBackendStop(c.BackendStop); err != nil { + return fail(err) + } + } + return &agentAttachment{ + connected: client.IsConnected, runs: runs, failed: failed, + close: func() error { + cancel() + client.Close() + // A subscription waits for the runs that are in flight, and the carrier + // is already closed, so nothing waits for it. + go func() { + for i := len(undo) - 1; i >= 0; i-- { + undo[i]() + } + }() + return nil + }, + }, nil +} + +func (c FollowConfig) tunnelAttacher(ctx context.Context, h worker.Handover) (worker.Attachment, error) { + if h.Creds == nil { + return nil, worker.ErrNoCredentials + } + runs := &atomic.Int32{} + work := NewWork() + consumer := counted{work, runs} + // No bus, so no broadcaster: the events of a run are published on the + // publisher of its delivery, which writes them to the frontend that drives + // the run. The cancel of a run is its request: the frontend ends the stream. + bridge := agents.NewEventBridge(nil, nil, "agent-worker-"+c.NodeID) + dispatcher := agents.NewNATSDispatcher(consumer, bridge, nil, c.APIURL, c.APIToken, 0) + // The consumers are registered before the tunnel starts, so that a frontend + // which connects at once finds a worker that can take its first run. + if err := dispatcher.Start(ctx); err != nil { + return nil, fmt.Errorf("starting dispatcher: %w", err) + } + if c.Jobs != nil { + if err := c.Jobs(ctx, consumer); err != nil { + _ = dispatcher.Stop() + return nil, err + } + } + cfg := Config{} + if c.Tool != nil { + cfg.MCPTool = Unary(c.Tool) + } + if c.Discovery != nil { + cfg.MCPDiscovery = Unary(c.Discovery) + } + if c.BackendStop != nil { + stop := c.BackendStop + cfg.BackendStop = func(_ context.Context, req workerctl.BackendStopRequest) error { + if req.Backend != "" { + stop(req.Backend) + } + return nil + } + } + rt, err := Start(ctx, Options{ + FrontendURL: c.FrontendURL, + NodeID: c.NodeID, + // The credential is new after every registration, so it is read at every + // dial and not once. + TunnelToken: h.Creds.TunnelToken, + Reauthorize: h.Creds.Reregister, + ControlToken: c.ControlToken, + Handler: Handler(cfg, work), + }) + if err != nil { + _ = dispatcher.Stop() + return nil, err + } + return &agentAttachment{ + connected: rt.Connected, runs: runs, + close: func() error { + err := rt.Close() + go func() { _ = dispatcher.Stop() }() + return err + }, + }, nil +} diff --git a/core/services/worker/config.go b/core/services/worker/config.go index 150a09542df8..41d8d189c2d2 100644 --- a/core/services/worker/config.go +++ b/core/services/worker/config.go @@ -89,7 +89,7 @@ type Config struct { FollowMaxDelay string `env:"LOCALAI_FOLLOW_MAX_DELAY" default:"10s" help:"Longest random wait before this worker attaches to the carrier that the cluster changed to" group:"registration" hidden:""` // NATS. The URL is needed only when the cluster runs on NATS. - NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Needed when the cluster runs on NATS; a cluster on the tunnel does not use it" group:"distributed"` + NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Optional: the frontend hands over the address when the cluster runs on NATS, and this overrides it" group:"distributed"` NatsJWT string `env:"LOCALAI_NATS_JWT" help:"NATS user JWT override (normally from registration nats_jwt)" group:"distributed"` NatsUserSeed string `env:"LOCALAI_NATS_USER_SEED" help:"NATS user signing seed override (normally from registration nats_user_seed)" group:"distributed"` NatsRequireAuth bool `env:"LOCALAI_NATS_REQUIRE_AUTH" default:"false" help:"Require NATS JWT+seed from registration or env" group:"distributed"` diff --git a/core/services/worker/follow_attach.go b/core/services/worker/follow_attach.go index 39fada8e4dd7..055beca3b101 100644 --- a/core/services/worker/follow_attach.go +++ b/core/services/worker/follow_attach.go @@ -38,6 +38,10 @@ type NATSLocal struct { Seed string Required bool TLS messaging.TLSFiles + // FallbackJWT and FallbackSeed are used when the frontend minted no + // credential, as an operator of an agent worker may configure. + FallbackJWT string + FallbackSeed string } // natsTLSFor returns the TLS files to use: the files of the operator, and the CA @@ -66,6 +70,11 @@ func ConnectNATS(ctx context.Context, local NATSLocal, h Handover, onFatal func( var opts []messaging.Option if h.Creds != nil && h.Creds.HasCredentials() { opts = append(opts, messaging.WithUserJWTProvider(h.Creds.Provider())) + } else if local.FallbackJWT != "" || local.FallbackSeed != "" { + if (local.FallbackJWT == "") != (local.FallbackSeed == "") { + return nil, errors.New("the fallback NATS JWT and seed must be set together") + } + opts = append(opts, messaging.WithUserJWT(local.FallbackJWT, local.FallbackSeed)) } else if local.Required { return nil, errors.New("NATS JWT+seed required: set LOCALAI_NATS_JWT/LOCALAI_NATS_USER_SEED or enable frontend minting") } From 277fcc96484118356a7d1fc4195009e967139fb3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 14:07:20 +0000 Subject: [PATCH 126/162] test(carrier): run the follower of a real worker through the changes of carrier The specs of the change no longer attach a fake worker by hand. A backend worker and an agent worker register through the real endpoints of the frontends, attach to the carrier the cluster names with the planes and the follower that Run uses, and follow each change by themselves from the answer of their heartbeat. The preflight reads what they report. From NATS to the tunnel and back with an install, a model that is loading on the backend, a held run and the renewals of a lease in flight: the workers attach to the target, keep NATS through the drain, release it when the cluster does, and nothing is lost or repeated. The backend answers over the new carrier without a restart. Two workers follow at once. A worker that restarts while a change is prepared follows it, and one that restarts after it starts on the carrier the cluster has, with no flag. An abort makes the worker let go of the target. A worker that predates carrier switching blocks a change and is left on NATS, unroutable and unreaped, by a forced one. A worker with no address is listed as unable to follow to NATS with the reason it reported, and stays on the tunnel after a forced change. The planes of a backend worker and its verbs are exported so that the specs can serve verbs of their own on the real carriers. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/agentworker/follow.go | 8 +- .../carrier/switch_integration_test.go | 746 +++++++++++++----- .../carrier/tunnel_integration_test.go | 12 + core/services/worker/follow_plane.go | 118 +++ core/services/worker/worker.go | 35 +- 5 files changed, 698 insertions(+), 221 deletions(-) create mode 100644 core/services/worker/follow_plane.go diff --git a/core/services/agentworker/follow.go b/core/services/agentworker/follow.go index b3fb4819d13b..147162ae2647 100644 --- a/core/services/agentworker/follow.go +++ b/core/services/agentworker/follow.go @@ -36,8 +36,8 @@ type FollowConfig struct { // BackendStop drops the sessions cached for a backend that goes away. BackendStop func(backend string) // Jobs registers the consumers of the queued jobs on the consumer of one - // carrier. It is called for each attachment. - Jobs func(ctx context.Context, consumer messaging.WorkConsumer) error + // carrier, which it is told. It is called for each attachment. + Jobs func(ctx context.Context, on cluster.Carrier, consumer messaging.WorkConsumer) error } // counted counts the runs that are in flight on one attachment. @@ -137,7 +137,7 @@ func (c FollowConfig) natsAttacher(ctx context.Context, h worker.Handover) (work } } if c.Jobs != nil { - if err := c.Jobs(actx, consumer); err != nil { + if err := c.Jobs(actx, cluster.CarrierNATS, consumer); err != nil { return fail(err) } } @@ -181,7 +181,7 @@ func (c FollowConfig) tunnelAttacher(ctx context.Context, h worker.Handover) (wo return nil, fmt.Errorf("starting dispatcher: %w", err) } if c.Jobs != nil { - if err := c.Jobs(ctx, consumer); err != nil { + if err := c.Jobs(ctx, cluster.CarrierTunnel, consumer); err != nil { _ = dispatcher.Stop() return nil, err } diff --git a/core/services/carrier/switch_integration_test.go b/core/services/carrier/switch_integration_test.go index c799e1bde190..9973e15597a4 100644 --- a/core/services/carrier/switch_integration_test.go +++ b/core/services/carrier/switch_integration_test.go @@ -1,11 +1,13 @@ package carrier_test import ( + "cmp" "context" "encoding/json" "errors" "fmt" "net" + "net/http" "net/http/httptest" "os" "runtime" @@ -21,12 +23,15 @@ import ( ggrpc "google.golang.org/grpc" "gorm.io/gorm" + "github.com/mudler/LocalAI/core/cli/workerregistry" clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/http/endpoints/localai" "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/agentworker" "github.com/mudler/LocalAI/core/services/carrier" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/jobs" + mcpRemote "github.com/mudler/LocalAI/core/services/mcp" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/testutil" @@ -34,6 +39,7 @@ import ( "github.com/mudler/LocalAI/core/services/worker" "github.com/mudler/LocalAI/core/services/workerctl" pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/natsauth" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -187,9 +193,17 @@ func (w rigWork) InFlight(ctx context.Context) (cluster.InFlight, error) { return cluster.InFlight{Jobs: int(n)}, err } -// switchRig is a cluster in one process: a database, a NATS server, a worker that -// can be reached over both carriers, an agent worker on each carrier, and -// frontend replicas that are built the way the application builds them. +// rigNATS is what the frontends of the rig hand over about NATS. +type rigNATS struct{ url string } + +func (r rigNATS) WorkerURL(context.Context) (string, error) { return r.url, nil } +func (r rigNATS) CAPEM() (string, error) { return "", nil } +func (r rigNATS) ClientTLS() bool { return false } + +// switchRig is a cluster in one process: a database, a NATS server, workers that +// run the code of a real worker to attach to either carrier, agent workers that +// do the same, and frontend replicas that are built the way the application +// builds them. type switchRig struct { ctx context.Context db *gorm.DB @@ -203,28 +217,280 @@ type switchRig struct { timings atomic.Pointer[cluster.Timings] - // The worker. + replicas []*switchReplica + + // onMove, when set, is called by the replica whose switch made a move of the + // row. A spec uses it to kill the leader between two moves. + onMove atomic.Pointer[func(r *switchReplica, row cluster.CarrierRow)] +} + +// workerProc is a backend worker: the follower and the planes of a real worker +// around a fake backend that serves gRPC. The backend outlives the worker, as a +// backend process does when its worker restarts. +type workerProc struct { + rig *switchRig + name string nodeID string backend *fakeBackend - control *controlPlane grpcAddr string httpAddr string + dir string + httpSrv *http.Server + fol *worker.Follower + cancel context.CancelFunc + + mu sync.Mutex + installs map[cluster.Carrier]int + renewals map[cluster.Carrier]int + installGate chan struct{} + installDone chan error +} - // What the worker answered on each carrier. - natsRenewals, natsInstalls atomic.Int32 - installGate chan struct{} - installDone chan error +type workerOpts struct { + // addr starts the worker with an address, which makes it dual-capable. + addr bool + // reuse starts a worker around the backend and the ports of one that stopped. + reuse *workerProc + // maxDelay is the longest random wait before the worker attaches. + maxDelay time.Duration +} - workerNATS *messaging.Client +func (w *workerProc) count(m map[cluster.Carrier]int, c cluster.Carrier) int { + w.mu.Lock() + defer w.mu.Unlock() + return m[c] +} +func (w *workerProc) installsOn(c cluster.Carrier) int { return w.count(w.installs, c) } +func (w *workerProc) renewalsOn(c cluster.Carrier) int { return w.count(w.renewals, c) } + +// verbs serves the verbs of a load on the carrier it is given, and counts them. +func (w *workerProc) verbs(v worker.VerbServer) error { + c := v.Carrier() + if err := v.HandleWithProgress(workerctl.VerbBackendInstall, func(_ context.Context, body []byte, progress func(workerctl.BackendInstallProgressEvent)) (any, error) { + var req workerctl.BackendInstallRequest + _ = json.Unmarshal(body, &req) + w.mu.Lock() + w.installs[c]++ + w.mu.Unlock() + progress(workerctl.BackendInstallProgressEvent{OpID: req.OpID, Percentage: 50}) + if req.OperationID == "load-slow" { + <-w.installGate + } + return workerctl.BackendInstallReply{Success: true, Address: w.grpcAddr, ProcessInstance: "i1", ReportsOperations: true}, nil + }); err != nil { + return err + } + if err := v.Handle(workerctl.VerbModelOp, func(_ context.Context, body []byte) (any, error) { + var req workerctl.OperationRequest + _ = json.Unmarshal(body, &req) + w.mu.Lock() + w.renewals[c]++ + w.mu.Unlock() + return workerctl.OperationReply{Renewed: req.Renew, Completed: req.Complete}, nil + }); err != nil { + return err + } + if err := v.Handle(workerctl.VerbModelStop, func(context.Context, []byte) (any, error) { + return workerctl.ModelStopReply{Matched: true, Terminated: true}, nil + }); err != nil { + return err + } + if err := v.Handle(workerctl.VerbBackendList, func(context.Context, []byte) (any, error) { + return workerctl.BackendListReply{}, nil + }); err != nil { + return err + } + return v.Handle(workerctl.VerbModelsRunning, func(context.Context, []byte) (any, error) { + return workerctl.ModelsRunningReply{}, nil + }) +} - // The agent worker. - agentID string +// stop ends the worker process: its follower and its HTTP server. The backend +// keeps serving. +func (w *workerProc) stop() { + w.cancel() + w.fol.Close() + if w.httpSrv != nil { + nodes.ShutdownFileTransferServer(w.httpSrv) + w.httpSrv = nil + } +} - replicas []*switchReplica +// startWorker registers a worker with a frontend and runs it the way Run does: +// the plane, the verbs, the attachment to the carrier the frontend names, and the +// follower. +func (rig *switchRig) startWorker(name string, o workerOpts) *workerProc { + GinkgoHelper() + w := &workerProc{rig: rig, name: name, installs: map[cluster.Carrier]int{}, renewals: map[cluster.Carrier]int{}, + installGate: make(chan struct{}), installDone: make(chan error, 1)} + if o.reuse != nil { + w.backend, w.grpcAddr, w.httpAddr, w.dir = o.reuse.backend, o.reuse.grpcAddr, o.reuse.httpAddr, o.reuse.dir + w.installGate, w.installDone = o.reuse.installGate, o.reuse.installDone + } else { + w.backend = &fakeBackend{loading: make(chan string, 1), gate: make(chan struct{})} + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + gs := ggrpc.NewServer() + pb.RegisterBackendServer(gs, w.backend) + go func() { _ = gs.Serve(lis) }() + DeferCleanup(gs.Stop) + w.grpcAddr = lis.Addr().String() + hl, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + w.httpAddr = hl.Addr().String() + Expect(hl.Close()).To(Succeed()) + w.dir = GinkgoT().TempDir() + } + ctx, cancel := context.WithCancel(rig.ctx) + w.cancel = cancel + front := rig.replicas[0].url + regClient := &workerregistry.RegistrationClient{FrontendURL: front, RegistrationToken: registrationToken} + + var routable atomic.Bool + routable.Store(o.addr) + body := func() map[string]any { + b := map[string]any{"name": name, "token": registrationToken, "routable": routable.Load(), "version": "switch-spec"} + if routable.Load() || rig.row().Active == cluster.CarrierNATS { + b["address"], b["http_address"] = w.grpcAddr, w.httpAddr + } + return b + } + credMgr := workerregistry.NewCredentialManager(func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + return regClient.RegisterFull(ctx, body()) + }, false) + res, err := credMgr.Acquire(ctx) + Expect(err).ToNot(HaveOccurred()) + w.nodeID = res.ID + boot := cluster.CarrierNATS + if res.Carrier == "tunnel" { + boot = cluster.CarrierTunnel + } else { + routable.Store(true) + } - // onMove, when set, is called by the replica whose switch made a move of the - // row. A spec uses it to kill the leader between two moves. - onMove atomic.Pointer[func(r *switchReplica, row cluster.CarrierRow)] + cfg := &worker.Config{RegisterTo: front, RegistrationToken: registrationToken, ServeAddr: w.grpcAddr} + if o.addr { + cfg.Addr = w.grpcAddr + } + plane := worker.NewBackendPlane(cfg, w.nodeID, w.httpAddr, worker.NATSLocal{}) + plane.Serve(w.verbs) + w.httpSrv, err = nodes.StartFileTransferServerWithControl(w.httpAddr, w.dir, w.dir, w.dir, registrationToken, 1<<30, nil, nil, plane.Control()) + Expect(err).ToNot(HaveOccurred()) + + delay := cmp.Or(o.maxDelay, 300*time.Millisecond) + w.fol = worker.NewFollower(worker.FollowerOptions{ + Attachers: plane.Attachers(), + Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { + return workerregistry.NewCredentialManagerFor(func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + b := body() + b["carrier"] = string(c) + return regClient.RegisterFull(ctx, b) + }, false, string(c)) + }, + Heartbeat: func(ctx context.Context, b map[string]any) (*workerregistry.HeartbeatReply, error) { + return regClient.HeartbeatFull(ctx, w.nodeID, b) + }, + Cannot: plane.Cannot(routable.Load()), + Interval: 300 * time.Millisecond, + MaxDelay: delay, + }) + _, err = w.fol.Attach(ctx, boot, credMgr, res) + Expect(err).ToNot(HaveOccurred()) + go w.fol.Run(ctx) + DeferCleanup(w.stop) + return w +} + +// agentProc is an agent worker around the follower of a real one. +type agentProc struct { + nodeID string + fol *worker.Follower + cancel context.CancelFunc +} + +func (a *agentProc) stop() { + a.cancel() + a.fol.Close() +} + +// startAgent registers an agent worker and runs it the way the agent worker +// command does. Every job it runs is recorded with the carrier it ran on. +func (rig *switchRig) startAgent(name string) *agentProc { + GinkgoHelper() + ctx, cancel := context.WithCancel(rig.ctx) + a := &agentProc{cancel: cancel} + front := rig.replicas[0].url + regClient := &workerregistry.RegistrationClient{FrontendURL: front, RegistrationToken: registrationToken} + body := func() map[string]any { + return map[string]any{"name": name, "node_type": "agent", "token": registrationToken, "version": "switch-spec"} + } + credMgr := workerregistry.NewCredentialManager(func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + return regClient.RegisterFull(ctx, body()) + }, false) + res, err := credMgr.Acquire(ctx) + Expect(err).ToNot(HaveOccurred()) + a.nodeID = res.ID + boot := cluster.CarrierNATS + if res.Carrier == "tunnel" { + boot = cluster.CarrierTunnel + } + serve := agentworker.FollowConfig{ + NodeID: a.nodeID, FrontendURL: front, ControlToken: registrationToken, + Subject: "agent.execute", Queue: "agent-workers", APIURL: "http://127.0.0.1:1", + Tool: func(context.Context, mcpRemote.MCPToolRequest) mcpRemote.MCPToolResponse { + return mcpRemote.MCPToolResponse{} + }, + Discovery: func(context.Context, mcpRemote.MCPDiscoveryRequest) mcpRemote.MCPDiscoveryResponse { + return mcpRemote.MCPDiscoveryResponse{} + }, + BackendStop: func(string) {}, + Jobs: func(ctx context.Context, on cluster.Carrier, consumer messaging.WorkConsumer) error { + _, err := consumer.Consume(ctx, messaging.WorkMCPCI, 0, rig.runs.handler(string(on))) + return err + }, + } + a.fol = worker.NewFollower(worker.FollowerOptions{ + Attachers: serve.Attachers(), + Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { + return workerregistry.NewCredentialManagerFor(func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + b := body() + b["carrier"] = string(c) + return regClient.RegisterFull(ctx, b) + }, false, string(c)) + }, + Heartbeat: func(ctx context.Context, b map[string]any) (*workerregistry.HeartbeatReply, error) { + return regClient.HeartbeatFull(ctx, a.nodeID, b) + }, + Cannot: serve.Cannot(), + Interval: 300 * time.Millisecond, + MaxDelay: 300 * time.Millisecond, + }) + _, err = a.fol.Attach(ctx, boot, credMgr, res) + Expect(err).ToNot(HaveOccurred()) + go a.fol.Run(ctx) + DeferCleanup(a.stop) + return a +} + +// attachedOf says where the frontends believe a node is attached, which is what +// the node reported in its last heartbeat. +func (rig *switchRig) attachedOf(nodeID string) []cluster.Carrier { + GinkgoHelper() + got, err := nodes.NewSwitchWorkers(rig.nodeReg, time.Hour).AttachedCarriers(rig.ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + return got +} + +// reports waits until a node has reported what it can follow. +func (rig *switchRig) reports(nodeID string) { + GinkgoHelper() + Eventually(func() string { + n, err := rig.nodeReg.Get(rig.ctx, nodeID) + if err != nil { + return "" + } + return n.Follow + }, "20s", "50ms").ShouldNot(BeEmpty(), "the node reports what it can follow") } type switchReplica struct { @@ -237,6 +503,7 @@ type switchReplica struct { cur atomic.Pointer[carrier.Set] bus *carrier.Broadcaster commands *carrier.Commands + clients *carrier.Clients queue *carrier.WorkQueue window *carrier.Window swapper *carrier.Swapper @@ -351,6 +618,14 @@ func (rig *switchRig) startReplica(id string) *switchReplica { e := echo.New() e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(rig.nodeReg, r.tunnels)) e.GET(tunnel.PeerPath, clusterapi.PeerHandler(rig.clusterR, r.sessions.Accept)) + // The routes a worker registers and heartbeats on, as the application mounts + // them. + nodeOpts := []localai.RegisterOption{ + localai.WithCarrierReader(rig.store), localai.WithTunnelDisconnector(r.tunnels), + localai.WithNATSHandover(rigNATS{url: rig.natsURL}), + } + e.POST("/api/node/register", localai.RegisterNodeEndpoint(rig.nodeReg, registrationToken, true, nil, "", natsauth.Config{}, nodeOpts...)) + e.POST("/api/node/:id/heartbeat", localai.HeartbeatEndpoint(rig.nodeReg, nodeOpts...)) srv := httptest.NewServer(e) DeferCleanup(srv.Close) r.url = srv.URL @@ -383,6 +658,8 @@ func (rig *switchRig) startReplica(id string) *switchReplica { }, workers.AgentsAttached) r.commands = carrier.NewCommands(&r.cur) r.commands.UseWindow(r.window) + r.clients = carrier.NewClients(&r.cur) + r.clients.UseWindow(r.window) r.queue = carrier.NewWorkQueue(&r.cur) var err2 error @@ -439,12 +716,6 @@ func (rig *switchRig) row() cluster.CarrierRow { return r } -// report says what the worker told the frontend about its carriers. -func (rig *switchRig) report(id string, attached []cluster.Carrier, follow []cluster.Carrier, followErr string) { - GinkgoHelper() - Expect(rig.nodeReg.SetCarrierReport(rig.ctx, id, nodes.CarrierReport{Attached: attached, Follow: follow, FollowError: followErr})).To(Succeed()) -} - func (rig *switchRig) newJob() string { GinkgoHelper() task := &jobs.TaskRecord{UserID: "u1", Name: "t", Model: "m", Enabled: true} @@ -513,8 +784,7 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { } ctx, cancel := context.WithCancel(context.Background()) DeferCleanup(cancel) - rig = &switchRig{ctx: ctx, natsURL: url, runs: newRunLog(), - installGate: make(chan struct{}), installDone: make(chan error, 1)} + rig = &switchRig{ctx: ctx, natsURL: url, runs: newRunLog()} rig.timings.Store(&cluster.Timings{PrepareTimeout: 20 * time.Second, TransitionWindow: 20 * time.Second, MaxDrain: 4 * time.Second}) rig.db, rig.dsn = testutil.SetupTestDBWithDSN() Expect(cluster.Migrate(ctx, rig.db)).To(Succeed()) @@ -528,112 +798,33 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Expect(err).ToNot(HaveOccurred()) _, _, err = rig.store.Seed(ctx, cluster.CarrierNATS, "seed") Expect(err).ToNot(HaveOccurred()) - - // The worker: a gRPC backend, an HTTP control plane for the tunnel, and - // responders on NATS for the same verbs. - rig.backend = &fakeBackend{loading: make(chan string, 1)} - lis, err := net.Listen("tcp", "127.0.0.1:0") - Expect(err).ToNot(HaveOccurred()) - gs := ggrpc.NewServer() - pb.RegisterBackendServer(gs, rig.backend) - go func() { _ = gs.Serve(lis) }() - DeferCleanup(gs.Stop) - rig.grpcAddr = lis.Addr().String() - rig.control = &controlPlane{address: rig.grpcAddr} - dir := GinkgoT().TempDir() - hl, err := net.Listen("tcp", "127.0.0.1:0") - Expect(err).ToNot(HaveOccurred()) - rig.httpAddr = hl.Addr().String() - Expect(hl.Close()).To(Succeed()) - httpSrv, err := nodes.StartFileTransferServerWithControl(rig.httpAddr, dir, dir, dir, registrationToken, 1<<30, nil, nil, rig.control.handler()) - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(func() { nodes.ShutdownFileTransferServer(httpSrv) }) - - node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken), Address: rig.grpcAddr, HTTPAddress: rig.httpAddr} - Expect(rig.nodeReg.Register(ctx, node, true)).To(Succeed()) - rig.nodeID = node.ID - Expect(rig.nodeReg.SetTunnelTokenHash(ctx, rig.nodeID, hashOf(workerToken))).To(Succeed()) - rig.report(rig.nodeID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") - - rig.workerNATS, err = messaging.New(url) - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(rig.workerNATS.Close) - reply := func(v any) []byte { b, _ := json.Marshal(v); return b } - _, err = rig.workerNATS.SubscribeReply(messaging.SubjectNodeBackendInstall(rig.nodeID), func(data []byte, respond func([]byte)) { - var req workerctl.BackendInstallRequest - _ = json.Unmarshal(data, &req) - rig.natsInstalls.Add(1) - if req.OperationID == "load-slow" { - <-rig.installGate - } - respond(reply(workerctl.BackendInstallReply{Success: true, Address: rig.grpcAddr, ProcessInstance: "i1", ReportsOperations: true})) - }) - Expect(err).ToNot(HaveOccurred()) - _, err = rig.workerNATS.SubscribeReply(messaging.SubjectNodeModelOp(rig.nodeID), func(data []byte, respond func([]byte)) { - var req workerctl.OperationRequest - _ = json.Unmarshal(data, &req) - rig.natsRenewals.Add(1) - respond(reply(workerctl.OperationReply{Renewed: req.Renew, Completed: req.Complete})) - }) - Expect(err).ToNot(HaveOccurred()) - - // The agent worker: a consumer on NATS. The claim queue's consumer is the - // dispatch loop of a replica, which hands work to an agent worker that - // holds a tunnel; see followAgent. - agent := &nodes.BackendNode{Name: "agent-1", NodeType: nodes.NodeTypeAgent, TokenHash: hashOf(registrationToken)} - Expect(rig.nodeReg.Register(ctx, agent, true)).To(Succeed()) - rig.agentID = agent.ID - Expect(rig.nodeReg.SetTunnelTokenHash(ctx, rig.agentID, hashOf(agentToken))).To(Succeed()) - rig.report(rig.agentID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") - agentNATS, err := messaging.New(url) - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(agentNATS.Close) - _, err = messaging.NewNATSWorkConsumer(agentNATS).Consume(ctx, messaging.WorkMCPCI, 0, rig.runs.handler("nats")) - Expect(err).ToNot(HaveOccurred()) }) - // followWorker connects the tunnel of the worker to a replica and tells the - // frontend that it is attached to both carriers, as a worker that follows a - // change does. - followWorker := func(to *switchReplica) { + settled := func(c cluster.Carrier) func() bool { + return func() bool { + r := rig.row() + return r.State == cluster.StateStable && r.Active == c + } + } + // change requests a change and waits until it has settled. + change := func(to cluster.Carrier, force bool) { GinkgoHelper() - t, err := worker.StartTunnel(rig.ctx, worker.TunnelConfig{ - FrontendURL: to.url, NodeID: rig.nodeID, Token: func() string { return workerToken }, - Services: map[string]worker.LocalService{ - tunnel.StreamTagGRPC: func(ctx context.Context, _ string) (net.Conn, error) { - var d net.Dialer - return d.DialContext(ctx, "tcp", rig.grpcAddr) - }, - tunnel.StreamTagHTTP: func(ctx context.Context, _ string) (net.Conn, error) { - var d net.Dialer - return d.DialContext(ctx, "tcp", rig.httpAddr) - }, - }, - }) + _, _, err := rig.request(to, force) Expect(err).ToNot(HaveOccurred()) - DeferCleanup(t.Close) - Eventually(func() bool { return to.tunnels.Holds(rig.nodeID) }, "20s", "50ms").Should(BeTrue()) - rig.report(rig.nodeID, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + rig.eventually(settled(to), "the change to "+string(to)+" settles") } - followAgent := func(to *switchReplica) *agentworker.Runtime { + drained := func() { GinkgoHelper() - work := agentworker.NewWork() - _, err := work.Consume(rig.ctx, messaging.WorkMCPCI, 0, rig.runs.handler("tunnel")) - Expect(err).ToNot(HaveOccurred()) - rt, err := agentworker.Start(rig.ctx, agentworker.Options{ - FrontendURL: to.url, NodeID: rig.agentID, TunnelToken: func() string { return agentToken }, - ControlToken: registrationToken, Handler: agentworker.Handler(agentworker.Config{}, work), - }) - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(rt.Close) - Eventually(func() bool { return to.tunnels.Holds(rig.agentID) }, "20s", "50ms").Should(BeTrue()) - rig.report(rig.agentID, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") - return rt + rig.eventually(func() bool { return rig.row().Draining == "" }, "the drain ends") } It("moves a cluster with work in flight from NATS to the tunnel and back, and loses and repeats nothing", func() { a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") Expect(a.active()).To(Equal(cluster.CarrierNATS)) + w := rig.startWorker("w1", workerOpts{addr: true}) + agent := rig.startAgent("agent-1") + rig.reports(w.nodeID) + rig.reports(agent.nodeID) // Chatter: both replicas publish a numbered message every few milliseconds // all through the changes, and every replica listens. @@ -673,7 +864,7 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { return case <-time.After(100 * time.Millisecond): } - _, err := a.commands.OperationControl(rig.nodeID, workerctl.OperationRequest{Renew: []string{"load-1"}}) + _, err := a.commands.OperationControl(w.nodeID, workerctl.OperationRequest{Renew: []string{"load-1"}}) renewMu.Lock() if err != nil { renewErrs = append(renewErrs, err) @@ -685,12 +876,22 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { } }) - // An install that is under way on NATS, and a run that is held on NATS. + // An install that is under way on NATS, a model that is loading on the + // backend, and a run that is held on NATS. go func() { - _, err := a.commands.InstallBackendOp(rig.nodeID, "llama-cpp", "m", "", 0, "op-slow", "load-slow", time.Minute, nil) - rig.installDone <- err + _, err := a.commands.InstallBackendOp(w.nodeID, "llama-cpp", "m", "", 0, "op-slow", "load-slow", time.Minute, nil) + w.installDone <- err + }() + Eventually(func() int { return w.installsOn(cluster.CarrierNATS) }, "10s").Should(Equal(1)) + loadDone := make(chan error, 1) + go func() { + res, err := a.clients.NewClient(w.nodeID, w.grpcAddr, false).LoadModel(rig.ctx, &pb.ModelOptions{Model: "slow"}) + if err == nil && !res.Success { + err = errors.New(res.Message) + } + loadDone <- err }() - Eventually(rig.natsInstalls.Load, "10s").Should(BeEquivalentTo(1)) + Eventually(w.backend.loading, "10s").Should(Receive(Equal("slow"))) heldJob := rig.newJob() release := rig.runs.hold(heldJob) rig.enqueue(a, heldJob) @@ -704,31 +905,33 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { } rig.eventually(rig.completed(onNATS...), "the jobs queued on NATS ran before the change") - // The tunnel side of the worker and of the agent worker is up, and the - // preflight, which is a dry run, passes. - followAgent(a) - followWorker(b) + // Nothing was attached by hand: the workers report what they can follow, and + // the preflight, which is a dry run, passes. rig.availability() report, err := a.sw.Preflight(rig.ctx, cluster.CarrierTunnel) Expect(err).ToNot(HaveOccurred()) Expect(report.OK).To(BeTrue(), "%+v", report.Blockers) Expect(report.Replicas).To(HaveLen(2)) + Expect(report.Workers).To(HaveLen(2)) Expect(report.InFlight.Jobs).To(BeNumerically(">=", 1)) - By("changing to the tunnel with the install, the run and the renewals in flight") + By("changing to the tunnel with the install, the load, the run and the renewals in flight") began := time.Now() _, _, err = rig.request(cluster.CarrierTunnel, false) Expect(err).ToNot(HaveOccurred()) - rig.eventually(func() bool { - r := rig.row() - return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel - }, "the change commits and settles") + rig.eventually(settled(cluster.CarrierTunnel), "the change commits and settles") GinkgoWriter.Printf("change to the tunnel took %s\n", time.Since(began)) rig.eventually(func() bool { return a.active() == cluster.CarrierTunnel && b.active() == cluster.CarrierTunnel }, "both replicas use the tunnel") Expect(rig.row().Draining).To(Equal(cluster.CarrierNATS)) + By("watching the workers follow, with no restart") + followed := time.Now() + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierTunnel)) + Eventually(func() []cluster.Carrier { return rig.attachedOf(agent.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierTunnel)) + GinkgoWriter.Printf("both workers reported the tunnel %s after the commit was seen\n", time.Since(followed)) + // New work goes to the claim queue, and the dispatch loop of a replica - // drives it on the agent worker that holds a tunnel. + // drives it on the agent worker, which follows by itself. var onTunnel []string for range 8 { id := rig.newJob() @@ -737,17 +940,19 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { } rig.eventually(rig.completed(onTunnel...), "the jobs queued after the change ran on the tunnel") - By("letting the install and the run finish where they began") - close(rig.installGate) + By("letting the install, the load and the run finish where they began") + close(w.installGate) var installErr error - Eventually(rig.installDone, "20s").Should(Receive(&installErr)) + Eventually(w.installDone, "20s").Should(Receive(&installErr)) Expect(installErr).ToNot(HaveOccurred(), "an install that began on NATS ends on NATS") + close(w.backend.gate) + Eventually(loadDone, "20s").Should(Receive(BeNil()), "the load of the model went on while the control moved: the backend never stopped") Expect(rig.status(heldJob)).To(Equal("running")) close(release) rig.eventually(rig.completed(heldJob), "the run that began on NATS reports its result, which the frontend still hears on NATS") By("waiting for the drain to end") - rig.eventually(func() bool { return rig.row().Draining == "" }, "the leader ends the drain") + drained() rig.eventually(func() bool { a.mu.Lock() defer a.mu.Unlock() @@ -758,6 +963,13 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { } return len(a.natsConns) > 0 }, "the old carrier is closed on a replica when the drain is over") + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel}), "the worker released NATS") + Eventually(func() []cluster.Carrier { return rig.attachedOf(agent.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + + By("using the backend that never restarted, over the tunnel") + out, err := b.clients.NewClient(w.nodeID, w.grpcAddr, false).Predict(rig.ctx, &pb.PredictOptions{Prompt: "after"}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out.Message)).To(Equal("echo: after")) close(renewStop) renewing.Wait() @@ -767,8 +979,8 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Expect(renewErrs).To(BeEmpty(), "no renewal failed across the change") Expect(maxGap).To(BeNumerically("<", 3*time.Second), "renewals were never missed for longer than the cadence of a lease") renewMu.Unlock() - Expect(rig.natsRenewals.Load()).To(BeNumerically(">", 0)) - Expect(rig.control.seenRenewals()).To(BeNumerically(">", 0), "renewals went to the worker over the tunnel once the tunnel was active") + Expect(w.renewalsOn(cluster.CarrierNATS)).To(BeNumerically(">", 0)) + Expect(w.renewalsOn(cluster.CarrierTunnel)).To(BeNumerically(">", 0), "renewals went to the worker over the tunnel once it was attached to it") // Every job ran once, on the carrier it was queued on. for _, id := range onNATS { @@ -820,32 +1032,63 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Expect(wrong).To(BeEmpty(), "of %d messages", total) }) - // toTunnel changes a cluster that runs on NATS to the tunnel and waits until it - // has settled. - toTunnel := func() { - GinkgoHelper() - _, _, err := rig.request(cluster.CarrierTunnel, false) + It("moves a worker to the tunnel and back with a model loading and a backend running, without a restart", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + w := rig.startWorker("w1", workerOpts{addr: true}) + rig.reports(w.nodeID) + loadDone := make(chan error, 1) + go func() { + res, err := a.clients.NewClient(w.nodeID, w.grpcAddr, false).LoadModel(rig.ctx, &pb.ModelOptions{Model: "slow"}) + if err == nil && !res.Success { + err = errors.New(res.Message) + } + loadDone <- err + }() + Eventually(w.backend.loading, "10s").Should(Receive(Equal("slow"))) + + followed := time.Now() + change(cluster.CarrierTunnel, false) + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "20ms").Should(ContainElement(cluster.CarrierTunnel)) + GinkgoWriter.Printf("the worker reported the tunnel %s after the change was requested\n", time.Since(followed)) + close(w.backend.gate) + Eventually(loadDone, "20s").Should(Receive(BeNil())) + + drained() + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + out, err := b.clients.NewClient(w.nodeID, w.grpcAddr, false).Predict(rig.ctx, &pb.PredictOptions{Prompt: "over the tunnel"}) Expect(err).ToNot(HaveOccurred()) - rig.eventually(func() bool { - r := rig.row() - return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel - }, "the change to the tunnel settles") - } - drained := func() { - GinkgoHelper() - rig.eventually(func() bool { return rig.row().Draining == "" }, "the drain ends") - } + Expect(string(out.Message)).To(Equal("echo: over the tunnel")) + reply, err := a.commands.InstallBackendOp(w.nodeID, "llama-cpp", "m", "", 0, "op-2", "load-2", time.Minute, nil) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + Expect(w.installsOn(cluster.CarrierTunnel)).To(Equal(1)) + + By("going back to NATS") + change(cluster.CarrierNATS, false) + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "20ms").Should(ContainElement(cluster.CarrierNATS)) + drained() + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierNATS})) + out, err = a.clients.NewClient(w.nodeID, w.grpcAddr, false).Predict(rig.ctx, &pb.PredictOptions{Prompt: "over NATS"}) + Expect(err).ToNot(HaveOccurred()) + Expect(string(out.Message)).To(Equal("echo: over NATS")) + _, err = b.commands.InstallBackendOp(w.nodeID, "llama-cpp", "m", "", 0, "op-3", "load-3", time.Minute, nil) + Expect(err).ToNot(HaveOccurred()) + Expect(w.installsOn(cluster.CarrierNATS)).To(Equal(1)) + Expect(w.backend.cancelledModels()).To(BeEmpty(), "no load was cancelled by the changes") + }) It("hands over the work that is still queued on the tunnel when the cluster goes back to NATS, and runs each unit once", func() { a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") - agent := followAgent(a) - followWorker(a) - toTunnel() + agent := rig.startAgent("agent-1") + w := rig.startWorker("w1", workerOpts{addr: true}) + rig.reports(agent.nodeID) + rig.reports(w.nodeID) + change(cluster.CarrierTunnel, false) drained() + Eventually(func() []cluster.Carrier { return rig.attachedOf(agent.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel})) // No agent worker holds a tunnel any more, so the units wait in the table. - Expect(agent.Close()).To(Succeed()) - rig.report(rig.agentID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS, cluster.CarrierTunnel}, "") + agent.stop() var queued []string for range 6 { id := rig.newJob() @@ -854,12 +1097,14 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { } Consistently(rig.runs.total, "1s", "100ms").Should(BeZero()) - _, _, err := rig.request(cluster.CarrierNATS, false) + // A consumer that stayed on NATS runs them after the hand-over. + stayed, err := messaging.New(rig.natsURL) Expect(err).ToNot(HaveOccurred()) - rig.eventually(func() bool { - r := rig.row() - return r.State == cluster.StateStable && r.Active == cluster.CarrierNATS - }, "the change back to NATS settles") + DeferCleanup(stayed.Close) + _, err = messaging.NewNATSWorkConsumer(stayed).Consume(rig.ctx, messaging.WorkMCPCI, 0, rig.runs.handler("nats")) + Expect(err).ToNot(HaveOccurred()) + + change(cluster.CarrierNATS, false) drained() rig.eventually(rig.completed(queued...), "the units that were queued on the tunnel run on NATS") @@ -874,6 +1119,63 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Expect(b.active()).To(Equal(cluster.CarrierNATS)) }) + It("moves two workers at once, and both serve on the new carrier", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + w1 := rig.startWorker("w1", workerOpts{addr: true}) + w2 := rig.startWorker("w2", workerOpts{addr: true}) + rig.reports(w1.nodeID) + rig.reports(w2.nodeID) + + change(cluster.CarrierTunnel, false) + for _, w := range []*workerProc{w1, w2} { + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierTunnel), w.name) + } + drained() + for _, w := range []*workerProc{w1, w2} { + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel}), w.name) + _, err := b.commands.InstallBackendOp(w.nodeID, "llama-cpp", "m", "", 0, "op-"+w.name, "load-"+w.name, time.Minute, nil) + Expect(err).ToNot(HaveOccurred(), w.name) + Expect(w.installsOn(cluster.CarrierTunnel)).To(Equal(1), w.name) + } + _ = a + }) + + It("starts a worker that restarts in the middle of a change on the carrier the cluster has then", func() { + a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + hold := make(chan struct{}) + b.mu.Lock() + b.holdBuild[cluster.CarrierTunnel] = hold + b.mu.Unlock() + w := rig.startWorker("w1", workerOpts{addr: true}) + rig.reports(w.nodeID) + + _, _, err := rig.request(cluster.CarrierTunnel, false) + Expect(err).ToNot(HaveOccurred()) + Eventually(func() cluster.State { return rig.row().State }, "10s").Should(Equal(cluster.StatePrepare)) + + By("restarting the worker while the change is prepared") + w.stop() + again := rig.startWorker("w1", workerOpts{addr: true, reuse: w}) + Expect(again.nodeID).To(Equal(w.nodeID)) + Expect(rig.attachedOf(again.nodeID)).To(Equal([]cluster.Carrier{cluster.CarrierNATS}), "the cluster is still on NATS") + Eventually(func() []cluster.Carrier { return rig.attachedOf(again.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierTunnel), "it follows the change that is being prepared") + + close(hold) + rig.eventually(settled(cluster.CarrierTunnel), "the change settles") + drained() + Eventually(func() []cluster.Carrier { return rig.attachedOf(again.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierTunnel})) + + By("restarting it once the cluster is on the tunnel") + again.stop() + third := rig.startWorker("w1", workerOpts{addr: true, reuse: again}) + Expect(rig.attachedOf(third.nodeID)).To(Equal([]cluster.Carrier{cluster.CarrierTunnel}), "it starts on the carrier the cluster has, with no flag") + Eventually(func() error { + _, err := a.commands.InstallBackendOp(third.nodeID, "llama-cpp", "m", "", 0, "op-r", "load-r", time.Minute, nil) + return err + }, "20s", "100ms").Should(Succeed(), "it serves on the tunnel as soon as the tunnel is up") + Expect(third.installsOn(cluster.CarrierTunnel)).To(Equal(1)) + }) + It("resolves a change when a replica dies in prepare, and the replica that starts again follows the row", func() { a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") c := rig.startReplica("replica-c") @@ -893,10 +1195,7 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { c.kill() close(hold) - rig.eventually(func() bool { - r := rig.row() - return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel - }, "the change commits without the replica that died") + rig.eventually(settled(cluster.CarrierTunnel), "the change commits without the replica that died") Expect(a.active()).To(Equal(cluster.CarrierTunnel)) Expect(b.active()).To(Equal(cluster.CarrierTunnel)) @@ -920,10 +1219,7 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { _, _, err := rig.request(cluster.CarrierTunnel, false) Expect(err).ToNot(HaveOccurred()) - rig.eventually(func() bool { - r := rig.row() - return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel - }, "the other replica settles the change from the row") + rig.eventually(settled(cluster.CarrierTunnel), "the other replica settles the change from the row") var survivors int for _, r := range []*switchReplica{a, b} { if r.active() == cluster.CarrierTunnel { @@ -935,6 +1231,8 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { It("aborts a change that is stuck in prepare, drops what was built for it, and leaves the cluster on its carrier", func() { a, b := rig.startReplica("replica-a"), rig.startReplica("replica-b") + w := rig.startWorker("w1", workerOpts{addr: true}) + rig.reports(w.nodeID) hold := make(chan struct{}) b.mu.Lock() b.holdBuild[cluster.CarrierTunnel] = hold @@ -944,6 +1242,7 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Expect(err).ToNot(HaveOccurred()) rig.eventually(func() bool { return a.buildCount(cluster.CarrierTunnel) == 1 }, "A built the target") Eventually(rig.listeners, "10s", "50ms").Should(Equal(int64(1)), "A listens on the database for the target") + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierTunnel), "the worker attaches to the target while it is prepared") row, err := a.sw.Abort(rig.ctx, "admin") Expect(err).ToNot(HaveOccurred()) @@ -954,9 +1253,14 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Eventually(rig.listeners, "20s", "100ms").Should(BeZero(), "what was built for the change is dropped on both replicas") Expect(a.active()).To(Equal(cluster.CarrierNATS)) Expect(b.active()).To(Equal(cluster.CarrierNATS)) + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierNATS}), "the worker lets go of the target when the change is aborted") // The cluster works as it did. + _, err = a.commands.OperationControl(w.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(err).ToNot(HaveOccurred()) id := rig.newJob() + w2 := rig.startAgent("agent-1") + rig.reports(w2.nodeID) rig.enqueue(a, id) rig.eventually(rig.completed(id), "work still runs on NATS") Expect(rig.runs.ranOn(id)).To(Equal([]string{"nats"})) @@ -965,44 +1269,108 @@ var _ = Describe("A change of carrier, end to end", Ordered, func() { Eventually(func() int { return b.chatter.get("after-abort") }, "10s").Should(Equal(heardBefore + 1)) }) - It("blocks a change on a worker that cannot follow, and a forced change leaves it unroutable and unreaped", func() { + It("blocks a change on a worker that predates carrier switching, and a forced change leaves it on NATS, unroutable and unreaped", func() { a, _ := rig.startReplica("replica-a"), rig.startReplica("replica-b") - rig.report(rig.nodeID, []cluster.Carrier{cluster.CarrierNATS}, []cluster.Carrier{cluster.CarrierNATS}, "it holds no tunnel credential") - Expect(rig.nodeReg.SetNodeModel(rig.ctx, rig.nodeID, "keep-me", 0, "loaded", rig.grpcAddr, 0)).To(Succeed()) + w := rig.startWorker("w1", workerOpts{addr: true}) + rig.reports(w.nodeID) + // An old worker: registered, alive, and silent about carriers. + old := &nodes.BackendNode{Name: "old", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken), Address: "127.0.0.1:1", HTTPAddress: "127.0.0.1:2"} + Expect(rig.nodeReg.Register(rig.ctx, old, true)).To(Succeed()) + Expect(rig.nodeReg.SetNodeModel(rig.ctx, old.ID, "keep-me", 0, "loaded", "127.0.0.1:1", 0)).To(Succeed()) + oldNATS, err := messaging.New(rig.natsURL) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(oldNATS.Close) + _, err = oldNATS.SubscribeReply(messaging.SubjectNodeModelOp(old.ID), func(_ []byte, respond func([]byte)) { + reply, _ := json.Marshal(workerctl.OperationReply{}) + respond(reply) + }) + Expect(err).ToNot(HaveOccurred()) - _, _, err := rig.request(cluster.CarrierTunnel, false) + _, _, err = rig.request(cluster.CarrierTunnel, false) var blocked *cluster.BlockedError Expect(errors.As(err, &blocked)).To(BeTrue(), "%v", err) Expect(blocked.Report.Blockers).To(ContainElement(SatisfyAll( HaveField("Kind", cluster.BlockerWorker), - HaveField("ID", rig.nodeID), - HaveField("Reason", ContainSubstring("it holds no tunnel credential")), + HaveField("ID", old.ID), + HaveField("Reason", ContainSubstring("predates carrier switching")), ))) + Expect(blocked.Report.Blockers).To(HaveLen(1), "the worker that follows is not a blocker") Expect(rig.row().State).To(Equal(cluster.StateStable)) Expect(rig.row().Epoch).To(Equal(int64(1))) _, _, err = rig.request(cluster.CarrierTunnel, true) Expect(err).ToNot(HaveOccurred()) - rig.eventually(func() bool { - r := rig.row() - return r.State == cluster.StateStable && r.Active == cluster.CarrierTunnel - }, "the forced change settles") + rig.eventually(settled(cluster.CarrierTunnel), "the forced change settles") Expect(rig.row().Force).To(BeTrue()) + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierTunnel)) - // While NATS drains the worker is still reached on it. - _, err = a.commands.OperationControl(rig.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) - Expect(err).ToNot(HaveOccurred(), "the window routes to the carrier the worker is attached to") + // While NATS drains the old worker is still reached on it. + _, err = a.commands.OperationControl(old.ID, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(errors.Is(err, nodes.ErrNoRoute)).To(BeFalse(), "the window routes to the carrier the worker is attached to: %v", err) drained() Eventually(func() bool { - _, err := a.commands.OperationControl(rig.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + _, err := a.commands.OperationControl(old.ID, workerctl.OperationRequest{Renew: []string{"op"}}) + return errors.Is(err, nodes.ErrNoRoute) + }, "20s", "100ms").Should(BeTrue(), "after the drain the old worker has no route, and that is all it has") + + models, err := rig.nodeReg.GetNodeModels(rig.ctx, old.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(HaveLen(1), "nothing was reaped") + node, err := rig.nodeReg.Get(rig.ctx, old.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Status).To(Equal(nodes.StatusHealthy)) + }) + + It("lists a worker that has no address as unable to follow to NATS, and leaves it on the tunnel, unroutable and unreaped, after a forced change", func() { + a, _ := rig.startReplica("replica-a"), rig.startReplica("replica-b") + w := rig.startWorker("w1", workerOpts{addr: true}) + rig.reports(w.nodeID) + change(cluster.CarrierTunnel, false) + drained() + + // A worker that holds a tunnel and has no address of its own. + only := rig.startWorker("tunnel-only", workerOpts{}) + rig.reports(only.nodeID) + Eventually(func() string { + n, err := rig.nodeReg.Get(rig.ctx, only.nodeID) + if err != nil { + return "" + } + return n.FollowError + }, "20s", "50ms").Should(ContainSubstring("--addr")) + Expect(rig.nodeReg.SetNodeModel(rig.ctx, only.nodeID, "keep-me", 0, "loaded", only.grpcAddr, 0)).To(Succeed()) + + _, _, err := rig.request(cluster.CarrierNATS, false) + var blocked *cluster.BlockedError + Expect(errors.As(err, &blocked)).To(BeTrue(), "%v", err) + Expect(blocked.Report.Blockers).To(ConsistOf(SatisfyAll( + HaveField("Kind", cluster.BlockerWorker), + HaveField("ID", only.nodeID), + HaveField("Reason", ContainSubstring("--addr")), + HaveField("Forceable", true), + ))) + + By("forcing the change") + change(cluster.CarrierNATS, true) + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(ContainElement(cluster.CarrierNATS), "the worker that can follow does") + drained() + Eventually(func() []cluster.Carrier { return rig.attachedOf(w.nodeID) }, "20s", "50ms").Should(Equal([]cluster.Carrier{cluster.CarrierNATS})) + Expect(rig.attachedOf(only.nodeID)).To(Equal([]cluster.Carrier{cluster.CarrierTunnel}), "the worker that cannot follow stays where it is") + + Eventually(func() bool { + _, err := a.commands.OperationControl(only.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) return errors.Is(err, nodes.ErrNoRoute) - }, "20s", "100ms").Should(BeTrue(), "after the drain the worker has no route, and that is all it has") + }, "20s", "100ms").Should(BeTrue(), "it has no route on NATS") + _, err = a.commands.OperationControl(w.nodeID, workerctl.OperationRequest{Renew: []string{"op"}}) + Expect(err).ToNot(HaveOccurred(), "the other worker is reached on NATS") - models, err := rig.nodeReg.GetNodeModels(rig.ctx, rig.nodeID) + models, err := rig.nodeReg.GetNodeModels(rig.ctx, only.nodeID) Expect(err).ToNot(HaveOccurred()) Expect(models).To(HaveLen(1), "nothing was reaped") - node, err := rig.nodeReg.Get(rig.ctx, rig.nodeID) + node, err := rig.nodeReg.Get(rig.ctx, only.nodeID) Expect(err).ToNot(HaveOccurred()) Expect(node.Status).To(Equal(nodes.StatusHealthy)) + Expect(node.FollowError).ToNot(BeEmpty(), "it says why") + Expect(only.fol.Attached()).To(Equal([]cluster.Carrier{cluster.CarrierTunnel}), "and it keeps its tunnel") }) }) diff --git a/core/services/carrier/tunnel_integration_test.go b/core/services/carrier/tunnel_integration_test.go index 4b40d279a721..76bbcc4a0f1b 100644 --- a/core/services/carrier/tunnel_integration_test.go +++ b/core/services/carrier/tunnel_integration_test.go @@ -50,6 +50,8 @@ type fakeBackend struct { loads []string cancelled []string loading chan string + // gate holds a load of the model "slow" until a spec closes it. + gate chan struct{} } func (b *fakeBackend) Health(context.Context, *pb.HealthMessage) (*pb.Reply, error) { @@ -63,6 +65,16 @@ func (b *fakeBackend) LoadModel(ctx context.Context, in *pb.ModelOptions) (*pb.R switch in.Model { case "fails": return &pb.Result{Success: false, Message: "out of memory"}, nil + case "slow": + select { + case b.loading <- in.Model: + default: + } + select { + case <-b.gate: + case <-ctx.Done(): + return nil, ctx.Err() + } case "hangs": select { case b.loading <- in.Model: diff --git a/core/services/worker/follow_plane.go b/core/services/worker/follow_plane.go new file mode 100644 index 000000000000..62eb3bc1c720 --- /dev/null +++ b/core/services/worker/follow_plane.go @@ -0,0 +1,118 @@ +package worker + +import ( + "context" + "net/http" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// VerbServer is a carrier on which a worker serves its control verbs. It is the +// exported face of the control server, for code that serves verbs of its own. +type VerbServer interface { + // Carrier says which carrier this server is on. + Carrier() cluster.Carrier + // Handle serves a verb with one reply. A nil reply is a verb without one. + Handle(verb string, h func(ctx context.Context, body []byte) (reply any, err error)) error + // HandleWithProgress serves a verb that reports progress while it runs. + HandleWithProgress(verb string, h func(ctx context.Context, body []byte, progress func(workerctl.BackendInstallProgressEvent)) (reply any, err error)) error +} + +type verbServer struct { + carrier cluster.Carrier + cs controlServer +} + +func (v verbServer) Carrier() cluster.Carrier { return v.carrier } + +func (v verbServer) Handle(verb string, h func(context.Context, []byte) (any, error)) error { + return v.cs.handle(controlVerb(verb), h) +} + +func (v verbServer) HandleWithProgress(verb string, h func(context.Context, []byte, func(workerctl.BackendInstallProgressEvent)) (any, error)) error { + return v.cs.handleWithProgress(controlVerb(verb), func(ctx context.Context, body []byte, p progressSink) (any, error) { + return h(ctx, body, p) + }) +} + +// BackendPlane is what a backend worker holds to attach to either carrier: the +// control plane on its HTTP server, the tunnel behind it, and the NATS +// connection with the verbs on the subjects of the node. Run builds one, and so +// can any code that runs a worker's carriers with verbs of its own. +type BackendPlane struct { + cfg *Config + nodeID string + local NATSLocal + tunnel *tunnelFollower + slot *nodes.ControlSlot + serve func(controlServer) error +} + +// NewBackendPlane returns the plane of a worker. httpAddr is where the HTTP +// server of the worker listens. +func NewBackendPlane(cfg *Config, nodeID, httpAddr string, local NATSLocal) *BackendPlane { + p := &BackendPlane{cfg: cfg, nodeID: nodeID, local: local} + p.tunnel = newTunnelFollower(tunnelFollowerOptions{Config: cfg, NodeID: nodeID, HTTPAddr: httpAddr}) + if p.tunnel.CanAttach() { + p.slot = &nodes.ControlSlot{} + p.tunnel.slot = p.slot + } + return p +} + +// Control is the handler to mount on the HTTP server of the worker under the +// control prefix. It is nil when the worker has no frontend URL a tunnel can use. +func (p *BackendPlane) Control() http.Handler { + if p.slot == nil { + return nil + } + return p.slot +} + +// CanTunnel reports whether the worker can attach to a tunnel at all. +func (p *BackendPlane) CanTunnel() bool { return p.tunnel.CanAttach() } + +// Serve sets the verbs of the worker. They are claimed on each carrier when the +// worker attaches to it. +func (p *BackendPlane) Serve(verbs func(VerbServer) error) { + p.serve = func(cs controlServer) error { + c := cluster.CarrierNATS + if _, ok := cs.(*httpControlServer); ok { + c = cluster.CarrierTunnel + } + return verbs(verbServer{carrier: c, cs: cs}) + } + p.tunnel.register = p.serve +} + +// setServe sets the verbs of the worker from inside the package. +func (p *BackendPlane) setServe(serve func(controlServer) error) { + p.serve = serve + p.tunnel.register = serve +} + +// Attachers returns the attacher of each carrier. +func (p *BackendPlane) Attachers() map[cluster.Carrier]Attacher { + return map[cluster.Carrier]Attacher{ + cluster.CarrierNATS: natsAttacher(p.local, p.nodeID, func(cs controlServer) error { + return p.serve(cs) + }), + cluster.CarrierTunnel: tunnelAttacher(p.tunnel), + } +} + +// Cannot says, for each carrier, why the worker cannot attach to it. routable is +// true when the frontends can dial the worker. +func (p *BackendPlane) Cannot(routable bool) map[cluster.Carrier]func(CarrierView) string { + return map[cluster.Carrier]func(CarrierView) string{ + cluster.CarrierNATS: NATSReasons(routable, p.local), + cluster.CarrierTunnel: func(CarrierView) string { + if !p.tunnel.CanAttach() { + return "this worker has no frontend URL that a tunnel can use" + } + return "" + }, + } +} diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 860b24ee33e5..cbc39a114730 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -8,7 +8,6 @@ import ( "fmt" "maps" "net" - "net/http" "os" "os/signal" "path/filepath" @@ -185,21 +184,12 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // answers as a path that does not exist, which is what a worker on NATS has // always answered. A worker that has no frontend URL a tunnel can use mounts // no plane at all. - follower := newTunnelFollower(tunnelFollowerOptions{ - Config: cfg, - NodeID: nodeID, - HTTPAddr: httpAddr, - }) - if onTunnel && !follower.CanAttach() { + plane := NewBackendPlane(cfg, nodeID, httpAddr, natsLocal) + if onTunnel && !plane.CanTunnel() { _, terr := tunnelEndpoint(cfg.RegisterTo, nodeID) return terr } - var controlMux http.Handler - if follower.CanAttach() { - controlSlot := &nodes.ControlSlot{} - follower.slot = controlSlot - controlMux = controlSlot - } + controlMux := plane.Control() httpServer, err := nodes.StartFileTransferServerWithControl(httpAddr, stagingDir, cfg.ModelsPath, dataDir, cfg.RegistrationToken, config.DefaultMaxUploadSize, readiness, ephemeralCapacity, controlMux, ml.BackendLogs()) if err != nil { return fmt.Errorf("starting HTTP file transfer server: %w", err) @@ -242,10 +232,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // The carriers of the worker. It follows the one the frontend names, so the // attachments are made by the same code at start and when the cluster changes. carriers := NewFollower(FollowerOptions{ - Attachers: map[cluster.Carrier]Attacher{ - cluster.CarrierNATS: natsAttacher(natsLocal, nodeID, func(cs controlServer) error { return registerVerbs(cs) }), - cluster.CarrierTunnel: tunnelAttacher(follower), - }, + Attachers: plane.Attachers(), Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { return workerregistry.NewCredentialManagerFor( registerFor(regClient, registrationBody, c), c == cluster.CarrierNATS && cfg.NatsAuthRequired() && !staticNATS, string(c)) @@ -253,16 +240,8 @@ func Run(ctx *cliContext.Context, cfg *Config) error { Heartbeat: func(ctx context.Context, body map[string]any) (*Beat, error) { return regClient.HeartbeatFull(ctx, nodeID, body) }, - Body: cfg.heartbeatBody, - Cannot: map[cluster.Carrier]func(CarrierView) string{ - cluster.CarrierNATS: NATSReasons(routable.Load(), natsLocal), - cluster.CarrierTunnel: func(CarrierView) string { - if !follower.CanAttach() { - return "this worker has no frontend URL that a tunnel can use" - } - return "" - }, - }, + Body: cfg.heartbeatBody, + Cannot: plane.Cannot(routable.Load()), Interval: heartbeatInterval, MaxDelay: cfg.followMaxDelay(), }) @@ -309,7 +288,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { } return nil } - follower.register = registerVerbs + plane.setServe(registerVerbs) // Attach to the carrier that the frontend named, with the credential that the // registration returned, and follow it from here on. From e2871134a449cbbaa6d8f36969e40a3785cf1afd Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 14:18:23 +0000 Subject: [PATCH 127/162] feat(nodes): demote a worker that is not on the active carrier once the cluster settles After a forced change, a worker that could not follow heartbeats over HTTP, which exists on both carriers, and cannot be reached. The scheduler kept choosing it, found no route, marked it unhealthy, and the next tick promoted it again, so a request that landed on it failed until the loop ended. The health monitor now reads what each worker reports. A worker that reports its carriers and is not attached to the active one is demoted, and it stays demoted until it reports the active carrier. Nothing is deleted: the status changes, its models stay, and its backend keeps running. Nothing is demoted while a change is under way or a carrier drains, because the window routes the workers then, and a worker that reports nothing is left to the tunnel read as before. A replica runs the check when a carrier is released, so the demotion does not wait for the next tick. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 16 ++++ core/services/nodes/health.go | 50 +++++++++++ core/services/nodes/health_presence_test.go | 92 +++++++++++++++++++++ 3 files changed, 158 insertions(+) diff --git a/core/application/distributed.go b/core/application/distributed.go index 547304f684f1..e1f12bcc3323 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -658,6 +658,22 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return set != nil && set.Name == cluster.CarrierTunnel }) + // A worker that reports it is not attached to the active carrier is demoted + // once the cluster has settled: after a forced change it is the worker that + // could not follow. While two carriers are attached the window routes it, so + // nothing is demoted then. The check also runs when a carrier is released, so + // that the demotion does not wait for the next tick. + healthMon.UseCarrier(func() (cluster.Carrier, bool) { + set := active.Load() + if set == nil { + return "", false + } + return set.Name, window.Previous() == nil + }) + broadcaster.OnReconnect(func() { + go healthMon.CheckNow(context.Background()) + }) + // The protocol of a change of carrier, and what this replica does in one. The // switch tells the other replicas of each move with a hint, so that they look // at the row at once. The hint is a courtesy; the poll is what counts. diff --git a/core/services/nodes/health.go b/core/services/nodes/health.go index 552a72f477ef..f7a412dec721 100644 --- a/core/services/nodes/health.go +++ b/core/services/nodes/health.go @@ -4,6 +4,7 @@ import ( "cmp" "context" "io" + "slices" "sync" "time" @@ -58,6 +59,36 @@ type HealthMonitor struct { presence NodePresenceReader reconnectGrace time.Duration tunnelActive func() bool + + // carrier is set by UseCarrier. It is written before the loop starts. + carrier func() (cluster.Carrier, bool) +} + +// UseCarrier makes the monitor demote a worker that reports it is not attached to +// the active carrier, once the cluster has settled. Such a worker heartbeats over +// HTTP, which exists on both carriers, and it cannot be reached: it could not +// follow a change, or it has not yet. The window of a change routes it through the +// previous carrier, so nothing is demoted while a change is under way or a carrier +// drains. +// +// active returns the active carrier and whether the cluster is settled: stable, +// with nothing draining. Call it before Start. +func (hm *HealthMonitor) UseCarrier(active func() (cluster.Carrier, bool)) { + hm.carrier = active +} + +// carrierDeparted reports whether a worker has told the deployment that it is not +// attached to the active carrier. A worker that reports nothing predates carrier +// switching and is not judged by this: the tunnel read covers it where it matters. +func (hm *HealthMonitor) carrierDeparted(node *BackendNode) bool { + if hm.carrier == nil || node == nil || node.Follow == "" { + return false + } + active, settled := hm.carrier() + if !settled { + return false + } + return !slices.Contains(splitCarriers(node.Attached), active) } // UsePresence makes the monitor read the tunnel of each node while the tunnel is @@ -150,6 +181,11 @@ func (hm *HealthMonitor) Stop() { } } +// CheckNow runs one round of checks at once. A replica calls it when a carrier is +// released, because the workers that did not follow are unreachable from that +// moment, and the next tick may be seconds away. +func (hm *HealthMonitor) CheckNow(ctx context.Context) { hm.checkAll(ctx) } + func (hm *HealthMonitor) run(ctx context.Context) { ticker := time.NewTicker(hm.checkInterval) defer ticker.Stop() @@ -250,6 +286,20 @@ func (hm *HealthMonitor) doCheckAll(ctx context.Context) { continue } + // The same holds for a worker that says it is not on the active carrier. It + // stays registered and it heartbeats, it is not scheduled, and nothing of it + // is deleted. It is promoted again when it reports the active carrier. + if hm.carrierDeparted(&node) { + if node.Status != StatusUnhealthy && node.Status != StatusOffline { + xlog.Warn("Node is heartbeating but is not attached to the active carrier; marking unhealthy", + "node", node.Name, "nodeID", node.ID, "attached", node.Attached, "followError", node.FollowError) + if err := hm.registry.MarkUnhealthy(ctx, node.ID); err != nil { + xlog.Error("Failed to mark a node that is off the active carrier unhealthy", "node", node.Name, "error", err) + } + } + continue + } + // Heartbeat is fresh — node is alive if node.Status == StatusUnhealthy || node.Status == StatusOffline { xlog.Info("Node recovered", "node", node.Name) diff --git a/core/services/nodes/health_presence_test.go b/core/services/nodes/health_presence_test.go index 60516bc10ded..1526a290b451 100644 --- a/core/services/nodes/health_presence_test.go +++ b/core/services/nodes/health_presence_test.go @@ -215,3 +215,95 @@ var _ = Describe("The health monitor and the connection rows of a real cluster", Expect(got.Status).To(Equal(StatusHealthy)) }) }) + +var _ = Describe("The health monitor and the carriers a worker reports", func() { + var ( + store *fakeNodeHealthStore + factory *fakeBackendClientFactory + active cluster.Carrier + settled bool + hm *HealthMonitor + ) + + BeforeEach(func() { + store = newFakeNodeHealthStore() + factory = newFakeBackendClientFactory() + active, settled = cluster.CarrierNATS, true + hm = newTestHealthMonitor(store, factory, true, 30*time.Second) + hm.perModelHealthCheck = true + hm.UseCarrier(func() (cluster.Carrier, bool) { return active, settled }) + }) + + node := func(id, status, attached, follow string) { + n := makeTestNode(id, id, "10.0.0.1:50051", status, freshTime()) + n.Attached, n.Follow = attached, follow + store.addNode(n) + store.addNodeModel(id, NodeModel{NodeID: id, ModelName: "m", Address: "127.0.0.1:50053"}) + } + + It("demotes a worker that reports it is not attached to the active carrier, and reaps nothing", func() { + node("n1", StatusHealthy, "tunnel", "tunnel") + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusUnhealthy)) + Expect(store.getCalls()).To(ContainElement("MarkUnhealthy:n1")) + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("MarkOffline"))) + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) + }) + + It("demotes it once and does not promote it while it stays off the active carrier", func() { + node("n1", StatusHealthy, "tunnel", "tunnel") + for range 4 { + hm.doCheckAll(context.Background()) + } + count := 0 + for _, c := range store.getCalls() { + if c == "MarkUnhealthy:n1" { + count++ + } + Expect(c).ToNot(Equal("MarkHealthy:n1")) + } + Expect(count).To(Equal(1)) + }) + + It("promotes it when it reports the active carrier", func() { + node("n1", StatusUnhealthy, "nats", "nats,tunnel") + hm.doCheckAll(context.Background()) + Expect(store.getCalls()).To(ContainElement("MarkHealthy:n1")) + }) + + It("leaves a worker that is attached to both carriers alone", func() { + node("n1", StatusHealthy, "nats,tunnel", "nats,tunnel") + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + }) + + It("leaves every worker alone while a change is under way or a carrier drains, because the window routes them", func() { + settled = false + node("n1", StatusHealthy, "tunnel", "tunnel") + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + }) + + It("leaves a worker that reports nothing alone: it predates carrier switching", func() { + node("n1", StatusHealthy, "", "") + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + }) + + It("does not probe the backends of a worker that it demoted", func() { + node("n1", StatusHealthy, "tunnel", "tunnel") + factory.setClient("127.0.0.1:50053", &fakeBackendClient{healthy: false}) + for range perModelMissThreshold + 1 { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).ToNot(ContainElement(ContainSubstring("RemoveNodeModel"))) + Expect(hm.misses).To(BeEmpty()) + }) + + It("does nothing when it was not told where to read the carrier", func() { + hm.carrier = nil + node("n1", StatusHealthy, "tunnel", "tunnel") + hm.doCheckAll(context.Background()) + Expect(store.getNode("n1").Status).To(Equal(StatusHealthy)) + }) +}) From 34032e754789c8d14470f83f018152d65ed6fd29 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 14:29:19 +0000 Subject: [PATCH 128/162] docs(agents): describe how workers follow the carrier The guide says how a worker learns the carrier from the answer of its heartbeat, how it gets the credential of the target and attaches with a random wait, what it reports, when it releases a carrier, why it can fail to follow, and how the health monitor treats a worker that does not follow. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 42 ++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 66cc296d22fc..6f2c610bf4e3 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -529,6 +529,48 @@ runner, where it fails so a runner without Docker cannot hide the check. The end-to-end specs in `tests/e2e/distributed` run the NATS implementations of the other seams against a real server, also through Docker. +## Workers follow the carrier + +A worker does not choose its carrier. Every heartbeat goes to the frontend over +HTTP, which exists on both carriers, and the answer carries the active carrier, +its epoch, the state of a change, the target and the carrier that drains +(`carrier`, `carrier_epoch`, `carrier_state`, `carrier_target`, +`carrier_draining`). `worker.Follower` reads it. The backend worker +(`worker.BackendPlane`) and the agent worker (`agentworker.FollowConfig`) give it +one `Attacher` for each carrier, so the code that attaches at start is the code +that attaches later. + +- At start the worker registers, learns the active carrier, and attaches to it. + It needs no NATS setting: the registration answer hands over the address + (`nats.worker_url`, else `nats.url`), the CA as PEM and the per-node JWT. A local + `--nats-url`, CA or credential still wins. A frontend that predates carriers + names none, and the worker then needs `--nats-url` as before. +- While a change is prepared the worker registers again with `carrier` set to the + target, which returns the credential of that carrier only. A tunnel token + replaces the previous one and ends the session of the replica, so a worker that + holds the tunnel never asks for one. Only the active carrier, the target of a + change and the carrier that drains can be asked for. +- It waits a random time of up to `LOCALAI_FOLLOW_MAX_DELAY` (ten seconds), then + attaches with a wait that grows from one to thirty seconds when the carrier + refuses. It keeps the old carrier open, and reports `attached`, `attached_epoch`, + `follow_capabilities` and `follow_error` in the heartbeat. That report is the + only source of where a worker is attached (`nodes.SwitchWorkers`). A worker that + reports nothing predates carrier switching and is on NATS. +- It closes a carrier when the cluster released it, the carrier it moved to is + connected, and no request runs on it (`inflight` in the control servers). It + reports first, then closes. Backend processes are not touched. A change that is + aborted makes it close the target. +- A worker that cannot follow keeps the carrier it has, reports why, and tries + again. A backend worker cannot follow to NATS without an address that the + frontends can dial (`--addr` or `--advertise-addr`; a worker that booted on NATS + counts, because its backends listen on every interface), or with a NATS server + that asks for a client certificate it does not have. Such a worker is a blocker + of the change, and a forced change leaves it where it is. +- After the cluster has settled, `HealthMonitor.UseCarrier` demotes a worker that + reports it is not attached to the active carrier. The change is status only, + like the tunnel read, and the worker is promoted when it reports the active + carrier. + ## Open items for a second carrier - `WorkHandler` returns only an error. A carrier whose stream handler must send a From e832c8d119b6696e8266f347ae2f02e5c8c6ecbd Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 14:29:41 +0000 Subject: [PATCH 129/162] fix(agentworker): tell the job registration which carrier it serves The job registration of an agent worker now receives the carrier it is called for, and the agent worker command accepts the new signature. A worker that fails to follow for one reason reports that reason alone, without the name of the carrier in front of it, because the preflight already names the carrier. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/agent_worker.go | 2 +- core/services/nodes/follow_report_test.go | 8 ++++---- core/services/worker/follow_carrier.go | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index b0c58e7dd807..fc4d533c0c8d 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -226,7 +226,7 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { }, // MCP CI jobs run here because the frontend cannot create MCP sessions // (for example stdio servers that use docker). - Jobs: func(ctx context.Context, consumer messaging.WorkConsumer) error { + Jobs: func(ctx context.Context, _ cluster.Carrier, consumer messaging.WorkConsumer) error { _, err := startMCPCIConsumer(ctx, consumer, apiURL, cmd.APIToken, mcpCIJobTimeout) return err }, diff --git a/core/services/nodes/follow_report_test.go b/core/services/nodes/follow_report_test.go index b4fafe286a03..04775033b43d 100644 --- a/core/services/nodes/follow_report_test.go +++ b/core/services/nodes/follow_report_test.go @@ -16,10 +16,10 @@ import ( var _ = Describe("The workers as the switch of the carrier sees them", func() { var ( - ctx context.Context - db *gorm.DB - reg *NodeRegistry - workers *SwitchWorkers + ctx context.Context + db *gorm.DB + reg *NodeRegistry + workers *SwitchWorkers ) register := func(name, nodeType, address string) *BackendNode { diff --git a/core/services/worker/follow_carrier.go b/core/services/worker/follow_carrier.go index 47e569e896c6..9cb330ef52c6 100644 --- a/core/services/worker/follow_carrier.go +++ b/core/services/worker/follow_carrier.go @@ -286,7 +286,7 @@ func (f *Follower) reportLocked() map[string]any { follow = append(follow, c) } case !held: - reasons = append(reasons, fmt.Sprintf("%s: %s", c, reason)) + reasons = append(reasons, reason) } } // A worker holds a carrier it can follow, even before it is connected, so a From 36574fbd9b33ef9da8d0e2773cfa13c7c14d5d79 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 14:32:30 +0000 Subject: [PATCH 130/162] fix(application): close the body of a heartbeat answer and mark the CA path as operator input The heartbeat client ignored the error of closing the body, and the read of the CA file for the handover of NATS takes a path that the operator configured, which the scanner reports as a file inclusion. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/carrier_runtime.go | 1 + core/cli/workerregistry/client.go | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go index 357576249302..ca93e88c3668 100644 --- a/core/application/carrier_runtime.go +++ b/core/application/carrier_runtime.go @@ -105,6 +105,7 @@ func (rt *carrierRuntime) CAPEM() (string, error) { if path == "" { return "", nil } + // #nosec G304 -- the operator set this path for the CA of the NATS server. raw, err := os.ReadFile(path) if err != nil { return "", fmt.Errorf("reading the NATS CA to hand to workers: %w", err) diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index e5409246aec3..a6712d5337e4 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -187,7 +187,7 @@ func (c *RegistrationClient) HeartbeatFull(ctx context.Context, nodeID string, b if err != nil { return nil, err } - defer resp.Body.Close() + defer func() { _ = resp.Body.Close() }() if resp.StatusCode < 200 || resp.StatusCode >= 300 { return nil, fmt.Errorf("heartbeat refused with status %d", resp.StatusCode) } From 2ea7ae3d72d890f1a56f9d571afec009aafda7c8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 14:59:37 +0000 Subject: [PATCH 131/162] test(worker): check what stops a backend worker from following to NATS The specs cover the missing address, the client certificate that cannot be handed over, a worker with no frontend URL for a tunnel, and the choice between the CA of the operator and the one the frontend hands over. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/worker/follow_carrier_test.go | 34 +++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/core/services/worker/follow_carrier_test.go b/core/services/worker/follow_carrier_test.go index 9f396cc8c1be..e0eef325eaf3 100644 --- a/core/services/worker/follow_carrier_test.go +++ b/core/services/worker/follow_carrier_test.go @@ -12,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/cli/workerregistry" "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/workerctl" . "github.com/onsi/ginkgo/v2" @@ -306,3 +307,36 @@ var _ = Describe("The requests that run on a control server", func() { Eventually(srv.InFlight, "5s").Should(BeZero()) }) }) + +var _ = Describe("What stops a backend worker from following to NATS", func() { + It("is a missing address that the frontends can dial", func() { + reason := NATSReasons(false, NATSLocal{})(CarrierView{}) + Expect(reason).To(ContainSubstring("--addr")) + Expect(NATSReasons(true, NATSLocal{})(CarrierView{})).To(BeEmpty()) + }) + + It("is a NATS server that asks for a client certificate which the worker lacks", func() { + view := CarrierView{NATSClientTLS: true} + Expect(NATSReasons(true, NATSLocal{})(view)).To(ContainSubstring("client certificate")) + Expect(NATSReasons(true, NATSLocal{TLS: messaging.TLSFiles{Cert: "c.pem"}})(view)).To(ContainSubstring("client certificate"), "a certificate without its key is not one") + Expect(NATSReasons(true, NATSLocal{TLS: messaging.TLSFiles{Cert: "c.pem", Key: "k.pem"}})(view)).To(BeEmpty()) + }) + + It("is nothing for the tunnel when the worker has a frontend URL, and a reason when it has none", func() { + with := NewBackendPlane(&Config{RegisterTo: "http://frontend:8080"}, "n1", "127.0.0.1:50050", NATSLocal{}) + Expect(with.Cannot(true)[cluster.CarrierTunnel](CarrierView{})).To(BeEmpty()) + Expect(with.Control()).ToNot(BeNil()) + Expect(with.CanTunnel()).To(BeTrue()) + + without := NewBackendPlane(&Config{}, "n1", "127.0.0.1:50050", NATSLocal{}) + Expect(without.Cannot(true)[cluster.CarrierTunnel](CarrierView{})).To(ContainSubstring("frontend URL")) + Expect(without.Control()).To(BeNil(), "a worker that cannot attach mounts no plane") + }) + + It("takes the CA that the frontend handed over only when the operator gave none", func() { + h := Handover{NATSCAPEM: "pem"} + Expect(natsTLSFor(messaging.TLSFiles{}, h).CAPEM).To(Equal([]byte("pem"))) + Expect(natsTLSFor(messaging.TLSFiles{CA: "ca.pem"}, h).CAPEM).To(BeEmpty()) + Expect(natsTLSFor(messaging.TLSFiles{}, Handover{}).Enabled()).To(BeFalse()) + }) +}) From 5fd1751e516666772ddb9e2fad29d18e0b7da9aa Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 15:04:38 +0000 Subject: [PATCH 132/162] fix(pgbus): offer a notification that arrives as the poll window ends The listener read the poll context before it offered the notification. When the 100 ms window ended between the read and the check, the message was taken off the wire and dropped with no log and no count. Under CPU load the spec that moves a cluster with work in flight lost a message in 4 of 8 runs. Offer the notification whenever the read returned one, whatever the context says. A spec now makes the connection return a notification at the moment the window ends. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/pgbus/listener.go | 62 ++++++++++++++++++----- core/services/pgbus/poll_internal_test.go | 52 +++++++++++++++++++ 2 files changed, 100 insertions(+), 14 deletions(-) create mode 100644 core/services/pgbus/poll_internal_test.go diff --git a/core/services/pgbus/listener.go b/core/services/pgbus/listener.go index 3ce42d005c48..4f3682bc051e 100644 --- a/core/services/pgbus/listener.go +++ b/core/services/pgbus/listener.go @@ -7,6 +7,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/mudler/xlog" ) @@ -240,29 +241,62 @@ func (b *Bus) listen(conn *pgx.Conn) { default: } - waitCtx, cancel := context.WithTimeout(b.ctx, listenPollInterval) - n, err := conn.WaitForNotification(waitCtx) - deadline := waitCtx.Err() != nil - cancel() - - switch { - case b.ctx.Err() != nil: + switch b.poll(conn.WaitForNotification) { + case pollStop: return - case deadline: - // The usual case: nothing arrived in the poll window. pgx keeps the - // connection when its read deadline passes, so this is not a failure. - continue - case err != nil: - xlog.Warn("Broadcast carrier lost its listen connection, dialling again", "error", err) + case pollRedial: replacement, ok := b.redial(conn) if !ok { return } conn = replacement - continue } + } +} + +// pollResult tells the loop of listen what to do after one wait. +type pollResult int + +const ( + pollContinue pollResult = iota + pollStop + pollRedial +) + +// waitFunc is the wait on the connection. It is a parameter so that a spec can +// make the connection return a notification at the moment the poll window ends. +type waitFunc func(ctx context.Context) (*pgconn.Notification, error) + +// poll waits one poll interval for a notification and offers it. +// +// A notification that WaitForNotification returned is offered whatever the poll +// context says afterwards. The context can expire between the read and the check, +// and a check that ran first would drop a message that was already off the wire, +// with no log and no count. +func (b *Bus) poll(wait waitFunc) pollResult { + waitCtx, cancel := context.WithTimeout(b.ctx, listenPollInterval) + n, err := wait(waitCtx) + deadline := waitCtx.Err() != nil + cancel() + + switch { + case err == nil && n != nil: b.offer(n.Channel, n.Payload) + if b.ctx.Err() != nil { + return pollStop + } + return pollContinue + case b.ctx.Err() != nil: + return pollStop + case deadline: + // The usual case: nothing arrived in the poll window. pgx keeps the + // connection when its read deadline passes, so this is not a failure. + return pollContinue + case err != nil: + xlog.Warn("Broadcast carrier lost its listen connection, dialling again", "error", err) + return pollRedial } + return pollContinue } // offer hands a notification to the resolver. It never blocks: the only job of diff --git a/core/services/pgbus/poll_internal_test.go b/core/services/pgbus/poll_internal_test.go new file mode 100644 index 000000000000..e11aac86c96b --- /dev/null +++ b/core/services/pgbus/poll_internal_test.go @@ -0,0 +1,52 @@ +package pgbus + +import ( + "context" + "errors" + + "github.com/jackc/pgx/v5/pgconn" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Internal on purpose: the window in which the poll context expires after the +// connection returned a notification is a few microseconds wide, and a spec that +// races for it passes by luck. These specs make the connection return at the +// moment the window ends. +var _ = Describe("one poll of the listen connection", func() { + var b *Bus + + BeforeEach(func() { + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + b = &Bus{ctx: ctx, inbound: make(chan inbound, 4)} + }) + + It("offers a notification that arrives together with an expired poll window", func() { + res := b.poll(func(ctx context.Context) (*pgconn.Notification, error) { + <-ctx.Done() // the window ends first ... + return &pgconn.Notification{Channel: "c", Payload: "p"}, nil // ... the read still won + }) + + Expect(res).To(Equal(pollContinue)) + Expect(b.inbound).To(Receive(Equal(inbound{channel: "c", payload: "p"}))) + }) + + It("keeps the connection when the window ends with nothing to read", func() { + res := b.poll(func(ctx context.Context) (*pgconn.Notification, error) { + <-ctx.Done() + return nil, ctx.Err() + }) + + Expect(res).To(Equal(pollContinue)) + Expect(b.inbound).ToNot(Receive()) + }) + + It("asks for a new connection when the read fails before the window ends", func() { + res := b.poll(func(context.Context) (*pgconn.Notification, error) { + return nil, errors.New("connection reset") + }) + + Expect(res).To(Equal(pollRedial)) + }) +}) From b9f038c1aba7a4e2b8d6721524f78b0b19905882 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 15:24:48 +0000 Subject: [PATCH 133/162] style(pgbus): format the poll spec Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/pgbus/poll_internal_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/services/pgbus/poll_internal_test.go b/core/services/pgbus/poll_internal_test.go index e11aac86c96b..75d20c7e9a41 100644 --- a/core/services/pgbus/poll_internal_test.go +++ b/core/services/pgbus/poll_internal_test.go @@ -24,7 +24,7 @@ var _ = Describe("one poll of the listen connection", func() { It("offers a notification that arrives together with an expired poll window", func() { res := b.poll(func(ctx context.Context) (*pgconn.Notification, error) { - <-ctx.Done() // the window ends first ... + <-ctx.Done() // the window ends first ... return &pgconn.Notification{Channel: "c", Payload: "p"}, nil // ... the read still won }) From c8489d415adc43ece57bb35de1947bcc357a42e3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 15:47:07 +0000 Subject: [PATCH 134/162] fix(cluster): make the NATS address of a change one value, and keep credentials out of it A change of carrier built the NATS carrier from the setting, which an admin could overwrite at any time. Two replicas could build from two servers during prepare and commit, and a replica that restarted after an edit sat disconnected. The setting also held whatever the flag held, including a user name, a password or a token in the query, and the status view only masked the user name. - The row of a change to NATS now carries the address, copied from the setting when the admin requests the change. While a change is under way every replica builds from the row, and a prepared set that was built from another address is dropped and built again. - The setting is written only when no change is under way and this replica reaches the server at the new address. A refused request stores nothing. - An address with a user name, a password, a query string or a fragment is refused on write. The flag is stored without them, and every view and log line drops them. A replica whose flag names the stored server still dials the flag, because that is where it keeps its credential. - A flag that names another server than the stored address stops the start and names both, with the way to move the broker. Before, the stored address won with one log line. A deployment that does not edit anything starts as before. - A replica with no flag no longer seeds the tunnel into a cluster that has NATS. The start fails when a worker or a replica is registered, or when the row says NATS. The address is stored before the row, so a replica that reads the row also finds the address. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 24 +++- core/application/carrier_runtime.go | 44 ++++-- core/application/distributed.go | 132 +++++++++++++++--- core/application/distributed_carrier_test.go | 81 ++++++++++- core/application/seed_carrier_test.go | 75 ++++++++++ .../http/endpoints/localai/cluster_carrier.go | 92 ++++++------ .../endpoints/localai/cluster_carrier_test.go | 54 +++++-- core/http/routes/cluster.go | 2 +- core/services/carrier/set.go | 4 + core/services/carrier/swap.go | 19 ++- core/services/carrier/swap_test.go | 37 ++++- core/services/cluster/carrier_store.go | 47 +++++-- core/services/cluster/natsurl.go | 72 ++++++++++ core/services/cluster/natsurl_test.go | 50 +++++++ core/services/cluster/settings.go | 4 +- core/services/cluster/switch.go | 17 ++- core/services/cluster/switch_test.go | 46 ++++++ 17 files changed, 686 insertions(+), 114 deletions(-) create mode 100644 core/application/seed_carrier_test.go create mode 100644 core/services/cluster/natsurl.go create mode 100644 core/services/cluster/natsurl_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 6f2c610bf4e3..c380b3362070 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -72,12 +72,28 @@ The row is seeded by the first replica that starts: NATS when it has a NATS URL (the flag `--nats-url`, or a URL stored in the cluster settings), the tunnel otherwise. `--nats-url` means what it always meant for a deployment on NATS, and it is no longer required: a deployment with only PostgreSQL runs on the tunnel. -The flag is copied into the cluster setting `nats.url` when none is stored, and a -stored URL wins afterwards, so that every replica uses the same one. A NATS URL + +The address is stored before the row, so a replica that finds a row for NATS also +finds the address. A replica with no flag and no stored address does not seed the +tunnel into a cluster that has NATS: when a worker or a replica is registered, or +the row names NATS, its start fails and says to set `--nats-url`. + +The flag is copied into the cluster setting `nats.url` when none is stored, without +its user name, password and query string. A flag that names another server than the +stored address stops the start and names both: the stored address is what every +replica uses, and the operator who moves a broker stores the new address +(`PUT /api/cluster/settings`) and restarts. A replica whose flag names the stored +server dials the flag, because that is where it keeps its own credential. A NATS URL that points at a server which is not up does not stop the start: `messaging.New` retries on a failed connect, so a frontend can start before its broker. A URL it cannot parse does. +The row of a change to NATS carries the address (`cluster_carrier.nats_url`), copied +from the setting when the admin requests the change. While the change is under +way every replica builds from the row, and a prepared set built from another +address is dropped and built again. The setting itself can be written only when no +change is under way and this replica reaches the server at the new address. + ## Changing the carrier An admin changes the carrier of a running cluster with `POST /api/cluster/carrier` @@ -162,7 +178,9 @@ of a change (`switch.prepare_timeout`, `switch.transition_window`, `switch.max_drain`; defaults 1m, 2m and 15m; the flags `--carrier-prepare-timeout`, `--carrier-transition-window` and `--carrier-max-drain` are the fallback). The runtime settings of the application are a file of one -process and cannot carry them. Credentials are never stored there. +process and cannot carry them. Credentials are never stored there: an address +with a user name, a password or a query string is refused on write, and every view +and log line drops them. ## The worker tunnel diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go index ca93e88c3668..daa7ffa33f28 100644 --- a/core/application/carrier_runtime.go +++ b/core/application/carrier_runtime.go @@ -19,7 +19,6 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/storage" "github.com/mudler/LocalAI/core/services/tunnel" - "github.com/mudler/LocalAI/pkg/sanitize" ) const ( @@ -136,16 +135,38 @@ func (rt *carrierRuntime) build(ctx context.Context, row cluster.CarrierRow, tar return nil, fmt.Errorf("%w: %q", cluster.ErrInvalidCarrier, target) } -func (rt *carrierRuntime) natsClient(ctx context.Context, wait time.Duration) (*messaging.Client, string, error) { - url, err := rt.natsURL(ctx) +// natsTarget is the address of the NATS server that a build uses. While a change +// is under way the row names it, so that every replica builds from the same +// address whatever happens to the setting; otherwise it is the setting, else the +// flag of this replica. +func (rt *carrierRuntime) natsTarget(ctx context.Context, row *cluster.CarrierRow) (string, error) { + if row != nil && row.State != cluster.StateStable && row.NATSURL != "" { + return row.NATSURL, nil + } + return rt.natsURL(ctx) +} + +// connectURL is the address to dial for target. The cluster stores the address +// without a credential. When the flag of this replica names the same server, the +// flag is dialled, because it is where this replica keeps its user name, +// password or token. +func (rt *carrierRuntime) connectURL(target string) string { + if flag := rt.cfg.Distributed.NatsURL; flag != "" && cluster.PublicNATSURL(flag) == target { + return flag + } + return target +} + +func (rt *carrierRuntime) natsClient(ctx context.Context, wait time.Duration, row *cluster.CarrierRow) (*messaging.Client, string, error) { + target, err := rt.natsTarget(ctx, row) if err != nil { return nil, "", err } - if url == "" { + if target == "" { return nil, "", errors.New("no NATS URL is configured: store one in the cluster settings (nats.url) or start the frontend with --nats-url") } - client, err := rt.natsClientAt(ctx, url, wait) - return client, url, err + client, err := rt.natsClientAt(ctx, rt.connectURL(target), wait) + return client, target, err } // natsClientAt connects to the NATS server at url with the credentials of this @@ -167,7 +188,7 @@ func (rt *carrierRuntime) natsClientAt(ctx context.Context, url string, wait tim for !client.IsConnected() { if time.Now().After(deadline) || ctx.Err() != nil { client.Close() - return nil, fmt.Errorf("the NATS server at %s did not answer within %s: check the address and the credentials", sanitize.URL(url), wait) + return nil, fmt.Errorf("the NATS server at %s did not answer within %s: check the address and the credentials", cluster.PublicNATSURL(url), wait) } select { case <-ctx.Done(): @@ -181,7 +202,7 @@ func (rt *carrierRuntime) natsClientAt(ctx context.Context, url string, wait tim // CheckNATS says whether this replica reaches the NATS server at url with its own // credentials. It opens a connection and closes it. func (rt *carrierRuntime) CheckNATS(ctx context.Context, url string) error { - client, err := rt.natsClientAt(ctx, url, natsProbeWait) + client, err := rt.natsClientAt(ctx, rt.connectURL(url), natsProbeWait) if err != nil { return err } @@ -194,7 +215,7 @@ func (rt *carrierRuntime) buildNATS(ctx context.Context, row cluster.CarrierRow, if wait { waitFor = natsConnectWait } - client, url, err := rt.natsClient(ctx, waitFor) + client, url, err := rt.natsClient(ctx, waitFor, &row) if err != nil { return nil, err } @@ -230,7 +251,8 @@ func (rt *carrierRuntime) buildNATS(ctx context.Context, row cluster.CarrierRow, client.Close() return nil, err } - xlog.Info("Connected to NATS", "url", sanitize.URL(url), "epoch", row.Epoch) + set.NATSURL = row.NATSURL + xlog.Info("Connected to NATS", "url", cluster.PublicNATSURL(url), "epoch", row.Epoch) return set, nil } @@ -277,7 +299,7 @@ func (rt *carrierRuntime) availability(ctx context.Context) map[cluster.Carrier] case cur != nil && cur.Name == cluster.CarrierNATS: out[cluster.CarrierNATS] = "" default: - client, _, err := rt.natsClient(ctx, natsProbeWait) + client, _, err := rt.natsClient(ctx, natsProbeWait, nil) if err != nil { out[cluster.CarrierNATS] = err.Error() } else { diff --git a/core/application/distributed.go b/core/application/distributed.go index e1f12bcc3323..a25b4358a20e 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -5,6 +5,7 @@ import ( "crypto/tls" "crypto/x509" "encoding/json" + "errors" "fmt" "io" "net" @@ -236,7 +237,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade if err != nil { return nil, fmt.Errorf("initializing cluster settings: %w", err) } - carrierRow, err := seedCarrier(context.Background(), carrierStore, settingsStore, cfg.Distributed.NatsURL, cfg.Distributed.InstanceID) + carrierRow, err := seedCarrier(context.Background(), authDB, carrierStore, settingsStore, cfg.Distributed.NatsURL, cfg.Distributed.InstanceID) if err != nil { return nil, err } @@ -683,6 +684,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade MaxDrain: cfg.Distributed.CarrierMaxDrain, } carrierSwitch, err := cluster.NewSwitch(cluster.SwitchOptions{ + NATSURL: runtime.natsURL, Store: carrierStore, Registry: clusterReg, Workers: workers, @@ -786,43 +788,137 @@ const leaderTick = 2 * time.Second // seedCarrier reads the cluster carrier row, and writes it when no replica has // yet. It returns the row as it is after the call. -func seedCarrier(ctx context.Context, store *cluster.CarrierStore, settings *cluster.SettingsStore, flagURL, instanceID string) (cluster.CarrierRow, error) { +// +// The NATS address of the cluster is the setting that is stored without a +// credential. The flag of this replica is copied there when none is stored. A +// flag that names another server than the stored address stops the start: the +// stored address is what every replica uses, so a replica that was given another +// one by mistake, or an operator who moved the broker and edited only the flag, +// would otherwise run on the old server and see a warning in a log. +// +// A deployment that had NATS before the row existed, and a cluster whose row says +// NATS, need an address. A replica that has none of it, and that sees no address +// stored, stops its start, because seeding the tunnel would move the whole +// cluster off the broker that its workers use. +// +// The address is stored before the row, so that a replica which reads the row and +// finds NATS also finds the address. +func seedCarrier(ctx context.Context, db *gorm.DB, store *cluster.CarrierStore, settings *cluster.SettingsStore, flagURL, instanceID string) (cluster.CarrierRow, error) { + // A flag that is not an address is not stored. The connection that is opened + // later says what is wrong with it. + flagPublic := "" + if flagURL != "" { + flagPublic = cluster.PublicNATSURL(flagURL) + if cluster.CheckNATSURL(flagPublic) != nil { + flagPublic = "" + } + } stored, _, err := settings.Get(ctx, cluster.SettingNATSURL) if err != nil { return cluster.CarrierRow{}, err } - seedWith := cluster.CarrierTunnel - if flagURL != "" || stored != "" { - seedWith = cluster.CarrierNATS + if flagPublic != "" && stored != "" && cluster.PublicNATSURL(stored) != flagPublic { + return cluster.CarrierRow{}, fmt.Errorf( + "the NATS address of this replica (--nats-url, %s) differs from the one stored for the cluster (%s). "+ + "Every replica uses the stored address. To move the cluster to the new address, store it with PUT /api/cluster/settings "+ + "({\"nats_url\": ...}) and restart; otherwise set --nats-url to the stored address or remove the flag", + flagPublic, cluster.PublicNATSURL(stored)) } - row, created, err := store.Seed(ctx, seedWith, instanceID) - if err != nil { + + row, err := store.Get(ctx) + switch { + case err == nil: + case errors.Is(err, cluster.ErrNotSeeded): + seedWith := cluster.CarrierTunnel + known := stored + if known == "" { + known = flagPublic + } + switch { + case known != "" || flagURL != "": + seedWith = cluster.CarrierNATS + default: + era, err := natsEraDeployment(ctx, db) + if err != nil { + return cluster.CarrierRow{}, err + } + if era { + return cluster.CarrierRow{}, errNoNATSURL("this cluster already has registered workers or replicas, which use NATS") + } + } + if stored == "" && flagPublic != "" { + if _, err := settings.SetIfAbsent(ctx, cluster.SettingNATSURL, flagPublic, instanceID); err != nil { + return cluster.CarrierRow{}, err + } + stored = flagPublic + } + var created bool + row, created, err = store.SeedWith(ctx, seedWith, instanceID, stored) + if err != nil { + return cluster.CarrierRow{}, fmt.Errorf("reading cluster carrier state: %w", err) + } + if created { + xlog.Info("Cluster carrier seeded", "carrier", row.Active, "epoch", row.Epoch) + } + default: return cluster.CarrierRow{}, fmt.Errorf("reading cluster carrier state: %w", err) } - if created { - xlog.Info("Cluster carrier seeded", "carrier", row.Active, "epoch", row.Epoch) + + // The address may have been stored by the replica that won the seed. + if stored == "" { + if stored, _, err = settings.Get(ctx, cluster.SettingNATSURL); err != nil { + return cluster.CarrierRow{}, err + } } - if flagURL != "" { + if stored == "" && flagPublic != "" { // The URL of the flag is copied into the cluster settings when none is - // stored, so that every replica uses the same one. A stored URL wins. - made, err := settings.SetIfAbsent(ctx, cluster.SettingNATSURL, flagURL, instanceID) + // stored, so that every replica uses the same one. + made, err := settings.SetIfAbsent(ctx, cluster.SettingNATSURL, flagPublic, instanceID) if err != nil { return cluster.CarrierRow{}, err } - switch { - case made: - xlog.Info("The NATS URL of this replica is now a setting of the cluster", "url", sanitize.URL(flagURL)) - case stored != "" && stored != flagURL: - xlog.Warn("The NATS URL of this replica differs from the one stored for the cluster; the stored one is used", - "flag", sanitize.URL(flagURL), "cluster", sanitize.URL(stored)) + if made { + xlog.Info("The NATS URL of this replica is now a setting of the cluster", "url", flagPublic) + stored = flagPublic } } + if stored == "" && flagURL == "" && (row.Active == cluster.CarrierNATS || row.Draining == cluster.CarrierNATS || row.Target == cluster.CarrierNATS) { + return cluster.CarrierRow{}, errNoNATSURL("the cluster carrier row says NATS") + } + if flagURL != "" && flagPublic != flagURL { + xlog.Warn("The NATS URL of this replica carries a credential or a query string. This replica keeps it for its own connection; the cluster stores the address without it", + "url", cluster.PublicNATSURL(flagURL)) + } if row.Active == cluster.CarrierTunnel && flagURL != "" { xlog.Info("The cluster runs on the tunnel carrier. The NATS URL of this replica is kept for a change back to NATS, and no NATS connection is opened") } return row, nil } +func errNoNATSURL(why string) error { + return fmt.Errorf("no NATS address is known: %s, so a replica without one would flip the cluster to the tunnel or sit disconnected. "+ + "Start this replica with --nats-url (LOCALAI_NATS_URL), or store the address with PUT /api/cluster/settings", why) +} + +// natsEraDeployment says whether the database shows a cluster that ran before it +// had a carrier row: a worker or a replica is registered. Such a cluster used +// NATS, because NATS was the only carrier. +func natsEraDeployment(ctx context.Context, db *gorm.DB) (bool, error) { + for _, table := range []string{"backend_nodes", "instances"} { + if !db.Migrator().HasTable(table) { + continue + } + var n int64 + if err := db.WithContext(ctx).Table(table).Limit(1).Count(&n).Error; err != nil { + return false, fmt.Errorf("looking for registered workers: %w", err) + } + if n > 0 { + return true, nil + } + } + return false, nil +} + // attachmentOf turns the carriers a worker is attached to into the routing form. func attachmentOf(cs []cluster.Carrier) carrier.Attachment { var a carrier.Attachment diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go index 8a2eb7820e0d..95009a73a71a 100644 --- a/core/application/distributed_carrier_test.go +++ b/core/application/distributed_carrier_test.go @@ -14,6 +14,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/testutil" "github.com/mudler/LocalAI/core/services/tunnel" "github.com/testcontainers/testcontainers-go" @@ -158,17 +159,85 @@ var _ = Describe("distributed startup and the cluster carrier", func() { Expect(url).To(Equal("nats://127.0.0.1:1")) }) - It("uses the NATS URL stored for the cluster and not the one of the flag", func() { + It("stops the start when the NATS URL of the flag differs from the one stored for the cluster, and names both", func() { store, err := cluster.NewSettingsStore(db) Expect(err).ToNot(HaveOccurred()) - Expect(store.Set(context.Background(), cluster.SettingNATSURL, "nats://[::1", "admin")).To(Succeed()) - cfg.Distributed.NatsURL = "nats://127.0.0.1:2" + Expect(store.Set(context.Background(), cluster.SettingNATSURL, "nats://old-broker:4222", "admin")).To(Succeed()) + cfg.Distributed.NatsURL = "nats://new-broker:4222" - _, err = initDistributed(cfg, db, nil, nil) - Expect(err).To(MatchError(ContainSubstring("connecting to NATS")), "the stored address, which cannot be parsed, was the one used") + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(HaveOccurred()) + Expect(svc).To(BeNil()) + Expect(err.Error()).To(ContainSubstring("nats://new-broker:4222")) + Expect(err.Error()).To(ContainSubstring("nats://old-broker:4222")) + Expect(err.Error()).To(ContainSubstring("/api/cluster/settings"), "it says how to resolve it") url, _, err := store.Get(context.Background(), cluster.SettingNATSURL) Expect(err).ToNot(HaveOccurred()) - Expect(url).To(Equal("nats://[::1"), "a flag does not overwrite a setting of the cluster") + Expect(url).To(Equal("nats://old-broker:4222"), "a flag does not overwrite a setting of the cluster") + }) + + It("starts on the stored NATS URL when the flag names the same server, and keeps the credential of the flag out of the database", func() { + cfg.Distributed.NatsURL = "nats://user:secret@127.0.0.1:1?token=abc" + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + url, ok, err := store.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://127.0.0.1:1")) + Expect(carrierRow().NATSURL).To(Equal("nats://127.0.0.1:1")) + + // A second start with the same flag is the same deployment. + cfg.Distributed.InstanceID = "replica-b" + second, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Shutdown) + }) + + It("does not seed the tunnel in a cluster that has workers, when the replica has no NATS URL", func() { + Expect(db.AutoMigrate(&nodes.BackendNode{})).To(Succeed()) + Expect(db.Create(&nodes.BackendNode{ID: "w1", Name: "w1", NodeType: "backend", Address: "w1:50051"}).Error).To(Succeed()) + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("--nats-url"))) + Expect(svc).To(BeNil()) + + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, err = store.Get(context.Background()) + Expect(err).To(MatchError(cluster.ErrNotSeeded), "the row is not written, so the next start with the flag seeds NATS") + }) + + It("does not start without a NATS URL when the row says NATS", func() { + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(context.Background(), cluster.CarrierNATS, "replica-b") + Expect(err).ToNot(HaveOccurred()) + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("the cluster carrier row says NATS"))) + Expect(svc).To(BeNil()) + }) + + It("starts without a NATS flag on the URL that the cluster stored", func() { + settings, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + Expect(settings.Set(context.Background(), cluster.SettingNATSURL, "nats://127.0.0.1:1", "admin")).To(Succeed()) + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(context.Background(), cluster.CarrierNATS, "replica-b") + Expect(err).ToNot(HaveOccurred()) + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierNATS)) }) Context("with a NATS server", func() { diff --git a/core/application/seed_carrier_test.go b/core/application/seed_carrier_test.go new file mode 100644 index 000000000000..b0da81bb386f --- /dev/null +++ b/core/application/seed_carrier_test.go @@ -0,0 +1,75 @@ +package application + +import ( + "context" + "fmt" + "sync" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("seeding the carrier of the cluster", func() { + It("lets a replica without a NATS flag start while replicas with one seed, whatever the order", func() { + for round := 0; round < 5; round++ { + db := testutil.SetupTestDB() + ctx := context.Background() + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + settings, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + + var wg sync.WaitGroup + errs := make([]error, 12) + start := make(chan struct{}) + for i := range errs { + flag := "" + if i%2 == 0 { + flag = "nats://broker:4222" + } + wg.Add(1) + go func() { + defer wg.Done() + <-start + _, errs[i] = seedCarrier(ctx, db, store, settings, flag, fmt.Sprintf("replica-%d", i)) + }() + } + close(start) + wg.Wait() + for i, err := range errs { + Expect(err).ToNot(HaveOccurred(), "replica %d", i) + } + + // Whoever won, the row and the address agree: a row that says NATS + // always has an address next to it. + row, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + if row.Active == cluster.CarrierNATS { + url, ok, err := settings.Get(ctx, cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://broker:4222")) + } + } + }) + + It("finds the address stored, and not the row alone, when a flag-less replica reads between the two writes of another", func() { + db := testutil.SetupTestDB() + ctx := context.Background() + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + settings, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + + // The first write of the replica with the flag has happened; the second has not. + _, err = settings.SetIfAbsent(ctx, cluster.SettingNATSURL, "nats://broker:4222", "replica-a") + Expect(err).ToNot(HaveOccurred()) + + row, err := seedCarrier(ctx, db, store, settings, "", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(row.Active).To(Equal(cluster.CarrierNATS), "an address is stored, so the cluster is on NATS and not on the tunnel") + Expect(row.NATSURL).To(Equal("nats://broker:4222")) + }) +}) diff --git a/core/http/endpoints/localai/cluster_carrier.go b/core/http/endpoints/localai/cluster_carrier.go index 49b2809dafaf..d866ec50d74b 100644 --- a/core/http/endpoints/localai/cluster_carrier.go +++ b/core/http/endpoints/localai/cluster_carrier.go @@ -5,15 +5,12 @@ import ( "errors" "fmt" "net/http" - "net/url" - "slices" "time" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/services/cluster" - "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/xlog" ) @@ -67,8 +64,8 @@ func carrierError(c echo.Context, status int, msg string, extra map[string]any) return c.JSON(status, body) } -// settingsView is the settings as an admin reads them. A password in the address -// of a server is masked. +// settingsView is the settings as an admin reads them. A user name, a password or a +// query string in the address of a server is dropped. func settingsView(all map[string]string) map[string]any { view := map[string]any{ "nats_url": "", @@ -78,10 +75,10 @@ func settingsView(all map[string]string) map[string]any { "max_drain": all[cluster.SettingMaxDrain], } if v := all[cluster.SettingNATSURL]; v != "" { - view["nats_url"] = sanitize.URL(v) + view["nats_url"] = cluster.PublicNATSURL(v) } if v := all[cluster.SettingNATSWorkerURL]; v != "" { - view["nats_worker_url"] = sanitize.URL(v) + view["nats_worker_url"] = cluster.PublicNATSURL(v) } return view } @@ -212,29 +209,29 @@ type settingsRequest struct { MaxDrain *string `json:"max_drain"` } -func validNATSURL(raw string) error { - u, err := url.Parse(raw) - if err != nil { - return fmt.Errorf("not a URL: %w", err) - } - if !slices.Contains([]string{"nats", "tls", "ws", "wss"}, u.Scheme) { - return fmt.Errorf("the scheme must be nats, tls, ws or wss, not %q", u.Scheme) - } - if u.Host == "" { - return errors.New("the address has no host") - } - return nil -} - // natsCheckTimeout bounds the reachability check of a saved URL. const natsCheckTimeout = 8 * time.Second -// PutClusterSettingsEndpoint stores settings of the cluster. Saving a NATS URL -// makes NATS available and checks that this replica reaches it; it does not -// switch the cluster to NATS. The change of carrier is a separate request with a -// dry run. Every replica checks the stored URL on its next report, and a dry run -// asks them to. -func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, prober ReplicaProber) echo.HandlerFunc { +// CarrierStatus reads the state of the cluster carrier. *cluster.Switch is one. +type CarrierStatus interface { + Status(ctx context.Context) (cluster.Report, error) +} + +// PutClusterSettingsEndpoint stores settings of the cluster. It stores nothing +// unless it can store everything: the request is checked as a whole first. +// +// A NATS URL is stored only when no change of carrier is under way and this +// replica reaches the server at that address. A change builds the NATS carrier +// from the address that the row of the change holds, so an edit during one +// would only confuse the operator; an address that does not answer would stay in +// the settings until a change to NATS failed on it. Saving an address makes NATS +// available; it does not switch the cluster to NATS. The change of carrier is a +// separate request with a dry run. Every replica checks the stored URL on its +// next report, and a dry run asks them to. +// +// No address may carry a user name, a password or a query string: the settings +// are shared, and the credentials of NATS stay on each replica. +func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, prober ReplicaProber, state CarrierStatus) echo.HandlerFunc { return func(c echo.Context) error { var req settingsRequest if err := c.Bind(&req); err != nil { @@ -261,13 +258,38 @@ func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, p continue } if ch.url { - if err := validNATSURL(*ch.val); err != nil { + if err := cluster.CheckNATSURL(*ch.val); err != nil { return carrierError(c, http.StatusBadRequest, fmt.Sprintf("%s: %v", ch.key, err), nil) } } else if _, err := time.ParseDuration(*ch.val); err != nil { return carrierError(c, http.StatusBadRequest, fmt.Sprintf("%s: %v", ch.key, err), nil) } } + + resp := map[string]any{"saved": true} + if req.NATSURL != nil { + st, err := state.Status(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + if st.State != cluster.StateStable { + return carrierError(c, http.StatusConflict, fmt.Sprintf( + "the NATS URL cannot be changed while a change of carrier is under way (%s to %s): wait for it to settle or abort it", + st.State, st.Row.Target), nil) + } + if *req.NATSURL != "" { + checkCtx, cancel := context.WithTimeout(ctx, natsCheckTimeout) + err := checker.CheckNATS(checkCtx, *req.NATSURL) + cancel() + if err != nil { + xlog.Warn("The NATS URL was refused: this replica cannot reach it", "error", err) + return carrierError(c, http.StatusUnprocessableEntity, fmt.Sprintf( + "nats_url: this replica cannot reach the NATS server at %s: %v", cluster.PublicNATSURL(*req.NATSURL), err), nil) + } + resp["nats"] = map[string]any{"reachable": true} + } + } + for _, ch := range changes { if ch.val == nil { continue @@ -279,7 +301,7 @@ func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, p err = settings.Set(ctx, ch.key, *ch.val, by) } switch { - case errors.Is(err, cluster.ErrUnknownSetting): + case errors.Is(err, cluster.ErrUnknownSetting), errors.Is(err, cluster.ErrCredentialInURL): return carrierError(c, http.StatusBadRequest, err.Error(), nil) case err != nil: return carrierError(c, http.StatusInternalServerError, err.Error(), nil) @@ -287,18 +309,6 @@ func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, p } xlog.Info("Cluster settings changed", "by", by) - resp := map[string]any{"saved": true} - if req.NATSURL != nil && *req.NATSURL != "" { - checkCtx, cancel := context.WithTimeout(ctx, natsCheckTimeout) - err := checker.CheckNATS(checkCtx, *req.NATSURL) - cancel() - result := map[string]any{"reachable": err == nil} - if err != nil { - result["reason"] = err.Error() - xlog.Warn("The NATS URL that was saved is not reachable from this replica", "error", err) - } - resp["nats"] = result - } if req.NATSURL != nil || req.NATSWorkerURL != nil { // The other replicas look at the new address and write what they find. prober.AskReplicas(ctx) diff --git a/core/http/endpoints/localai/cluster_carrier_test.go b/core/http/endpoints/localai/cluster_carrier_test.go index 32018eabf24a..3518ed2643ad 100644 --- a/core/http/endpoints/localai/cluster_carrier_test.go +++ b/core/http/endpoints/localai/cluster_carrier_test.go @@ -147,7 +147,7 @@ var _ = Describe("The admin API of the carrier", func() { report: cluster.Report{Active: cluster.CarrierNATS, Epoch: 4, State: cluster.StateStable, OK: true, Replicas: []cluster.ReplicaStatus{{ID: "a", Version: "v1"}}}, row: cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 5, State: cluster.StatePrepare, Target: cluster.CarrierTunnel}, } - settings = &memorySettings{m: map[string]string{cluster.SettingNATSURL: "nats://user:secret@broker:4222"}} + settings = &memorySettings{m: map[string]string{cluster.SettingNATSURL: "nats://user:secret@broker:4222?token=abc"}} nats = &fakeNATS{} }) @@ -163,6 +163,7 @@ var _ = Describe("The admin API of the carrier", func() { Expect(shown).To(ContainSubstring("broker:4222")) Expect(shown).ToNot(ContainSubstring("secret")) Expect(shown).ToNot(ContainSubstring("user")) + Expect(shown).ToNot(ContainSubstring("abc"), "a token in the query is not shown either") Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("secret"), "the stored value is not changed") }) }) @@ -236,7 +237,7 @@ var _ = Describe("The admin API of the carrier", func() { Describe("the settings", func() { It("stores a NATS URL, checks that this replica reaches it, and does not switch", func() { - code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"nats_url":"nats://broker2:4222"}`) + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"nats_url":"nats://broker2:4222"}`) Expect(code).To(Equal(http.StatusOK)) Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://broker2:4222")) Expect(nats.seen).To(Equal([]string{"nats://broker2:4222"})) @@ -244,36 +245,60 @@ var _ = Describe("The admin API of the carrier", func() { Expect(sw.requests).To(BeEmpty()) }) - It("still stores a URL that this replica cannot reach, and says why", func() { + It("refuses a URL that this replica cannot reach, says why, and stores nothing", func() { nats.err = errors.New("no answer within 3s") - code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"nats_url":"nats://broker2:4222"}`) - Expect(code).To(Equal(http.StatusOK)) - Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://broker2:4222")) - nb := body["nats"].(map[string]any) - Expect(nb["reachable"]).To(BeFalse()) - Expect(nb["reason"]).To(ContainSubstring("no answer")) + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, + `{"nats_url":"nats://broker2:4222","max_drain":"20m"}`) + Expect(code).To(Equal(http.StatusUnprocessableEntity)) + Expect(body["error"]).To(ContainSubstring("no answer")) + Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://user:secret@broker:4222?token=abc")) + Expect(settings.m).ToNot(HaveKey(cluster.SettingMaxDrain), "a refused request stores none of its parts") + }) + + It("refuses a NATS URL while a change of carrier is under way, and stores nothing", func() { + for _, st := range []cluster.State{cluster.StatePrepare, cluster.StateCommit} { + sw.report.State = st + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, + `{"nats_url":"nats://broker2:4222"}`) + Expect(code).To(Equal(http.StatusConflict), string(st)) + Expect(body["error"]).To(ContainSubstring("change of carrier")) + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("broker:4222")) + Expect(nats.seen).To(BeEmpty(), "the server is not even asked") + } + }) + + It("refuses an address that carries a credential or a query string", func() { + for _, bad := range []string{ + `{"nats_url":"nats://user:pass@broker2:4222"}`, + `{"nats_url":"nats://broker2:4222?token=abc"}`, + `{"nats_worker_url":"nats://user@broker2:4222"}`, + } { + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, bad) + Expect(code).To(Equal(http.StatusBadRequest), bad) + Expect(body["error"]).To(ContainSubstring("credentials stay on each replica"), bad) + } }) It("refuses a URL that is not a NATS address", func() { for _, bad := range []string{`{"nats_url":"http://broker:4222"}`, `{"nats_url":"nats://"}`, `{"nats_url":"::"}`} { - code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, bad) + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, bad) Expect(code).To(Equal(http.StatusBadRequest), bad) } - Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://user:secret@broker:4222")) + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("broker:4222")) }) It("clears a URL with an empty string", func() { - code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"nats_url":""}`) + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"nats_url":""}`) Expect(code).To(Equal(http.StatusOK)) Expect(settings.m).ToNot(HaveKey(cluster.SettingNATSURL)) }) It("stores the waits of a change, and refuses one that is not a positive duration", func() { - code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"max_drain":"20m","prepare_timeout":"90s"}`) + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"max_drain":"20m","prepare_timeout":"90s"}`) Expect(code).To(Equal(http.StatusOK)) Expect(settings.m[cluster.SettingMaxDrain]).To(Equal("20m")) Expect(settings.m[cluster.SettingPrepareTimeout]).To(Equal("90s")) - code, _ = call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}), http.MethodPut, `{"max_drain":"soon"}`) + code, _ = call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"max_drain":"soon"}`) Expect(code).To(Equal(http.StatusBadRequest)) }) @@ -282,6 +307,7 @@ var _ = Describe("The admin API of the carrier", func() { Expect(code).To(Equal(http.StatusOK)) Expect(body["nats_url"]).To(ContainSubstring("broker:4222")) Expect(body["nats_url"]).ToNot(ContainSubstring("secret")) + Expect(body["nats_url"]).ToNot(ContainSubstring("token")) }) }) diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go index 8d1ebb7b2fe8..bb39c177c2fc 100644 --- a/core/http/routes/cluster.go +++ b/core/http/routes/cluster.go @@ -67,5 +67,5 @@ func RegisterClusterAdminRoutes(e *echo.Echo, adminMw echo.MiddlewareFunc, admin e.GET(CarrierPath, localai.GetCarrierEndpoint(admin.Switch, admin.Settings), adminMw) e.POST(CarrierPath, localai.SwitchCarrierEndpoint(admin.Switch, admin.Prober), adminMw) e.GET(ClusterSettingsPath, localai.GetClusterSettingsEndpoint(admin.Settings), adminMw) - e.PUT(ClusterSettingsPath, localai.PutClusterSettingsEndpoint(admin.Settings, admin.NATS, admin.Prober), adminMw) + e.PUT(ClusterSettingsPath, localai.PutClusterSettingsEndpoint(admin.Settings, admin.NATS, admin.Prober, admin.Switch), adminMw) } diff --git a/core/services/carrier/set.go b/core/services/carrier/set.go index bccf21ffcebe..f9e0f0f33cf0 100644 --- a/core/services/carrier/set.go +++ b/core/services/carrier/set.go @@ -38,6 +38,10 @@ type Set struct { Name cluster.Carrier // Epoch is the epoch of the cluster row the set was built for. Epoch int64 + // NATSURL is the address of the NATS server that a NATS set was built from, + // as the row of the cluster named it. A set that was built from another + // address is not the set that the row asks for. Empty for the tunnel. + NATSURL string // Broadcaster carries fan-out and nothing else. A subject outside the // broadcast roots (see messaging.ValidateBroadcastSubject) is not served diff --git a/core/services/carrier/swap.go b/core/services/carrier/swap.go index 1afc6affbcb8..45ee91f29f05 100644 --- a/core/services/carrier/swap.go +++ b/core/services/carrier/swap.go @@ -308,11 +308,20 @@ func (s *Swapper) report(ctx context.Context, epoch int64, reason string) error // makes every subscription listen on it. func (s *Swapper) attach(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*Set, error) { switch { - case s.prepared != nil && s.prepared.Name == target: + case s.prepared != nil && s.prepared.Name == target && builtFrom(s.prepared, row): return s.prepared, nil - case s.draining != nil && s.draining.Name == target: + case s.draining != nil && s.draining.Name == target && builtFrom(s.draining, row): return s.draining, nil } + if s.prepared != nil && s.prepared.Name == target { + // Built from another address than the row names now. Every replica must + // build from the address of the row, so this set is dropped. + stale := s.prepared + s.prepared = nil + if stale != s.draining { + closeSet(s.dropListening(stale)) + } + } set, err := s.o.Build(ctx, row, target) if err != nil { return nil, err @@ -327,6 +336,12 @@ func (s *Swapper) attach(ctx context.Context, row cluster.CarrierRow, target clu return set, nil } +// builtFrom says whether set was built from the address that the row names. A +// row with no address, or a set of the other carrier, has nothing to compare. +func builtFrom(set *Set, row cluster.CarrierRow) bool { + return set.Name != cluster.CarrierNATS || row.NATSURL == "" || set.NATSURL == row.NATSURL +} + func (s *Swapper) prepare(ctx context.Context, row cluster.CarrierRow) error { set, err := s.attach(ctx, row, row.Target) if err != nil { diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go index b88756d4c21f..4152af932288 100644 --- a/core/services/carrier/swap_test.go +++ b/core/services/carrier/swap_test.go @@ -72,7 +72,11 @@ func (n *replicaNode) build(_ context.Context, row cluster.CarrierRow, target cl return nil, err } n.builds[target]++ - return n.newSet(target, row.Epoch), nil + set := n.newSet(target, row.Epoch) + if target == cluster.CarrierNATS { + set.NATSURL = row.NATSURL + } + return set, nil } // newSet builds a fake set over the shared bus of the carrier. The caller holds @@ -244,6 +248,37 @@ var _ = Describe("The swap of a replica", func() { }) Describe("prepare", func() { + It("builds again when the row names another NATS address than the set it holds was built from", func() { + a := join("a") + settle(cluster.CarrierTunnel) + row0 := row() + _, err := store.Transition(ctx, row0.Epoch, cluster.Change{ + Active: cluster.CarrierTunnel, State: cluster.StatePrepare, Target: cluster.CarrierNATS, + Draining: row0.Draining, DrainingUntil: row0.DrainingUntil, + By: "admin", NATSURL: "nats://broker-one:4222", + }) + Expect(err).ToNot(HaveOccurred()) + poll() + Expect(a.buildCount(cluster.CarrierNATS)).To(Equal(1)) + first := a.fake(cluster.CarrierNATS) + + // Same address, new epoch: the set is reused. + Expect(a.swapper.Reconcile(ctx, row())).To(Succeed()) + Expect(a.buildCount(cluster.CarrierNATS)).To(Equal(1)) + + // Another replica saw another address: the row decides. + _, err = store.Transition(ctx, row().Epoch, cluster.Change{ + Active: cluster.CarrierTunnel, State: cluster.StatePrepare, Target: cluster.CarrierNATS, + Draining: row0.Draining, DrainingUntil: row0.DrainingUntil, + By: "admin", NATSURL: "nats://broker-two:4222", + }) + Expect(err).ToNot(HaveOccurred()) + poll() + Expect(a.buildCount(cluster.CarrierNATS)).To(Equal(2), "a set built from the old address is not reused") + Expect(a.fake(cluster.CarrierNATS)).ToNot(BeIdenticalTo(first)) + Expect(a.life.seen()).To(ContainElement("close:nats"), "the set built from the old address is closed") + }) + It("builds the target while the old carrier still serves, listens on both, and reports ready", func() { a, b := join("a"), join("b") var heardOnNew atomic.Int32 diff --git a/core/services/cluster/carrier_store.go b/core/services/cluster/carrier_store.go index fb10c043bc2d..116d782e8726 100644 --- a/core/services/cluster/carrier_store.go +++ b/core/services/cluster/carrier_store.go @@ -79,6 +79,13 @@ type CarrierRow struct { // Note says in a sentence how the last change ended, or why it did. It is // for operators. Note string `gorm:"size:512;not null;default:''" json:"note,omitempty"` + // NATSURL is the address of the NATS server that this change, or the last + // change to NATS, builds from. The admin starts a change and the replicas + // build in it, and an edit of the setting in between must not give two + // replicas two servers, so every replica builds from this copy while a change + // is under way. It never carries a credential. Empty when the cluster has no + // NATS address. + NATSURL string `gorm:"column:nats_url;size:2048;not null;default:''" json:"nats_url,omitempty"` } func (CarrierRow) TableName() string { return "cluster_carrier" } @@ -94,6 +101,10 @@ type Change struct { By string Force bool Note string + // NATSURL replaces the address of the row when it is not empty. An empty + // value keeps the address that the row has, because only the start of a + // change to NATS knows which address to use. + NATSURL string } func (c Change) validate() error { @@ -149,6 +160,11 @@ func (s *CarrierStore) Get(ctx context.Context) (CarrierRow, error) { // whose insert won. Replicas that start together can all call it: the first // insert wins and the others read that row. func (s *CarrierStore) Seed(ctx context.Context, active Carrier, by string) (row CarrierRow, created bool, err error) { + return s.SeedWith(ctx, active, by, "") +} + +// SeedWith is Seed for a row that also carries the NATS address of the cluster. +func (s *CarrierStore) SeedWith(ctx context.Context, active Carrier, by, natsURL string) (row CarrierRow, created bool, err error) { if !active.valid() { return CarrierRow{}, false, fmt.Errorf("%w: %q", ErrInvalidCarrier, active) } @@ -159,6 +175,7 @@ func (s *CarrierStore) Seed(ctx context.Context, active Carrier, by string) (row State: StateStable, ChangedBy: by, ChangedAt: time.Now().UTC(), + NATSURL: natsURL, }) if res.Error != nil { return CarrierRow{}, false, fmt.Errorf("seeding cluster carrier: %w", res.Error) @@ -176,21 +193,25 @@ func (s *CarrierStore) Transition(ctx context.Context, from int64, change Change } // RETURNING hands back the row this update wrote, not a later one. var row CarrierRow + fields := map[string]any{ + "active": change.Active, + "state": change.State, + "target": change.Target, + "draining": change.Draining, + "draining_until": change.DrainingUntil, + "changed_by": change.By, + "changed_at": gorm.Expr("now()"), + "force": change.Force, + "note": change.Note, + "prev_epoch": gorm.Expr("epoch"), + "epoch": gorm.Expr("epoch + 1"), + } + if change.NATSURL != "" { + fields["nats_url"] = change.NATSURL + } res := s.db.WithContext(ctx).Model(&row).Clauses(clause.Returning{}). Where("id = ? AND epoch = ?", carrierRowID, from). - Updates(map[string]any{ - "active": change.Active, - "state": change.State, - "target": change.Target, - "draining": change.Draining, - "draining_until": change.DrainingUntil, - "changed_by": change.By, - "changed_at": gorm.Expr("now()"), - "force": change.Force, - "note": change.Note, - "prev_epoch": gorm.Expr("epoch"), - "epoch": gorm.Expr("epoch + 1"), - }) + Updates(fields) if res.Error != nil { return CarrierRow{}, fmt.Errorf("changing cluster carrier: %w", res.Error) } diff --git a/core/services/cluster/natsurl.go b/core/services/cluster/natsurl.go new file mode 100644 index 000000000000..1f6a0bf5d999 --- /dev/null +++ b/core/services/cluster/natsurl.go @@ -0,0 +1,72 @@ +package cluster + +import ( + "errors" + "fmt" + "net/url" + "slices" + "strings" +) + +// ErrCredentialInURL means an address carries a user name, a password or a query +// string. The cluster settings are shared and readable by every admin, so no +// credential is stored in them: a replica keeps its keys and tokens in its own +// flags and files. +var ErrCredentialInURL = errors.New("the address must not carry a user name, a password or a query string: " + + "credentials stay on each replica (flags, environment or files)") + +// CheckNATSURL says whether raw is an address that may be stored for the +// cluster: a NATS, TLS or WebSocket address with a host, and with no user name, +// no password, no query string and no fragment. +// +// The NATS client accepts several servers in one string, separated by commas. +// Each one is checked. +func CheckNATSURL(raw string) error { + for _, one := range strings.Split(raw, ",") { + if err := checkOneNATSURL(strings.TrimSpace(one)); err != nil { + return err + } + } + return nil +} + +func checkOneNATSURL(raw string) error { + u, err := url.Parse(raw) + if err != nil { + return fmt.Errorf("not a URL: %w", err) + } + if !slices.Contains([]string{"nats", "tls", "ws", "wss"}, u.Scheme) { + return fmt.Errorf("the scheme must be nats, tls, ws or wss, not %q", u.Scheme) + } + if u.Host == "" { + return errors.New("the address has no host") + } + if u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" { + return ErrCredentialInURL + } + return nil +} + +// PublicNATSURL returns raw without its user name, password, query string and +// fragment. Use it for every view and log line of an address that may come from +// a flag. It returns "***" for text that is not a URL. +func PublicNATSURL(raw string) string { + parts := strings.Split(raw, ",") + for i, one := range parts { + parts[i] = publicOneNATSURL(strings.TrimSpace(one)) + } + return strings.Join(parts, ",") +} + +func publicOneNATSURL(raw string) string { + u, err := url.Parse(raw) + if err != nil { + return "***" + } + u.User = nil + u.RawQuery = "" + u.ForceQuery = false + u.Fragment = "" + u.RawFragment = "" + return u.String() +} diff --git a/core/services/cluster/natsurl_test.go b/core/services/cluster/natsurl_test.go new file mode 100644 index 000000000000..a47094063b3b --- /dev/null +++ b/core/services/cluster/natsurl_test.go @@ -0,0 +1,50 @@ +package cluster_test + +import ( + "context" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("the NATS address of the cluster", func() { + It("accepts a plain address and a list of plain addresses", func() { + Expect(cluster.CheckNATSURL("nats://broker:4222")).To(Succeed()) + Expect(cluster.CheckNATSURL("tls://a:4222, tls://b:4222")).To(Succeed()) + }) + + It("refuses a user name, a password, a token in the query and a fragment", func() { + for _, bad := range []string{ + "nats://user@broker:4222", + "nats://user:pass@broker:4222", + "nats://broker:4222?token=abc", + "nats://broker:4222?", + "nats://broker:4222#frag", + "nats://a:4222,nats://u:p@b:4222", + } { + Expect(cluster.CheckNATSURL(bad)).To(MatchError(cluster.ErrCredentialInURL), bad) + } + }) + + It("shows an address without anything a credential could hide in", func() { + Expect(cluster.PublicNATSURL("nats://user:pass@broker:4222?token=abc#x")).To(Equal("nats://broker:4222")) + Expect(cluster.PublicNATSURL("nats://u:p@a:4222, nats://b:4222?t=1")).To(Equal("nats://a:4222,nats://b:4222")) + Expect(cluster.PublicNATSURL("")).To(Equal("")) + }) + + It("is not stored with a credential, by Set or by SetIfAbsent", func() { + store, err := cluster.NewSettingsStore(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + ctx := context.Background() + for _, key := range []string{cluster.SettingNATSURL, cluster.SettingNATSWorkerURL} { + Expect(store.Set(ctx, key, "nats://u:p@broker:4222", "x")).To(MatchError(cluster.ErrCredentialInURL)) + _, err := store.SetIfAbsent(ctx, key, "nats://broker:4222?token=1", "x") + Expect(err).To(MatchError(cluster.ErrCredentialInURL)) + _, ok, err := store.Get(ctx, key) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeFalse()) + } + }) +}) diff --git a/core/services/cluster/settings.go b/core/services/cluster/settings.go index af6223759a25..1be5b681c28c 100644 --- a/core/services/cluster/settings.go +++ b/core/services/cluster/settings.go @@ -37,8 +37,8 @@ const ( var ErrUnknownSetting = errors.New("unknown cluster setting") var knownSettings = map[string]func(string) error{ - SettingNATSURL: nil, - SettingNATSWorkerURL: nil, + SettingNATSURL: CheckNATSURL, + SettingNATSWorkerURL: CheckNATSURL, SettingPrepareTimeout: durationSetting, SettingTransitionWindow: durationSetting, SettingMaxDrain: durationSetting, diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index 2198899afd01..5bc43a416fc2 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -193,6 +193,11 @@ type SwitchOptions struct { // Timings returns the waits of a change, read at each use so that a change of // the settings applies to the next decision. Timings func() Timings + // NATSURL returns the address of the NATS server that a change to NATS builds + // from. Request copies it into the row, so that every replica builds from the + // same address whatever happens to the setting during the change. It may be + // nil. + NATSURL func(ctx context.Context) (string, error) // Liveness is the window of replica liveness. InstanceLiveness when zero. Liveness time.Duration // AvailabilityMaxAge is how old a replica's report of what it can build may @@ -388,11 +393,19 @@ func (s *Switch) Request(ctx context.Context, req Request) (CarrierRow, Report, if req.Force && len(report.Blockers) > 0 { note = fmt.Sprintf("forced past %d blocker(s)", len(report.Blockers)) } - next, err := s.transition(ctx, row.Epoch, Change{ + change := Change{ Active: row.Active, State: StatePrepare, Target: req.Target, Draining: row.Draining, DrainingUntil: row.DrainingUntil, Force: req.Force, Note: note, By: req.By, - }) + } + if req.Target == CarrierNATS && s.o.NATSURL != nil { + addr, err := s.o.NATSURL(ctx) + if err != nil { + return CarrierRow{}, report, fmt.Errorf("reading the NATS address of the cluster: %w", err) + } + change.NATSURL = PublicNATSURL(addr) + } + next, err := s.transition(ctx, row.Epoch, change) if errors.Is(err, ErrStaleEpoch) { return CarrierRow{}, report, ErrBusy } diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go index c13a68344e63..06587cf3f704 100644 --- a/core/services/cluster/switch_test.go +++ b/core/services/cluster/switch_test.go @@ -247,6 +247,52 @@ var _ = Describe("The switch of the carrier", func() { Expect(got.Force).To(BeFalse()) }) + Describe("the NATS address of a change", func() { + var current atomic.Pointer[string] + + BeforeEach(func() { + u := "nats://broker-one:4222" + current.Store(&u) + var err error + sw, err = cluster.NewSwitch(cluster.SwitchOptions{ + Store: store, Registry: reg, Workers: workers, Work: work, + Timings: func() cluster.Timings { return timings }, + NATSURL: func(context.Context) (string, error) { return *current.Load(), nil }, + }) + Expect(err).ToNot(HaveOccurred()) + // The cluster runs on the tunnel, so a change to NATS can start. + _, err = store.Transition(ctx, row().Epoch, cluster.Change{Active: cluster.CarrierTunnel, State: cluster.StateStable, By: "x"}) + Expect(err).ToNot(HaveOccurred()) + }) + + It("copies the address into the row, and keeps it through the commit and the settle", func() { + replica("a", "", "") + got, _, err := request(cluster.CarrierNATS, false) + Expect(err).ToNot(HaveOccurred()) + Expect(got.NATSURL).To(Equal("nats://broker-one:4222")) + + // The setting changes during the change. The row does not. + u := "nats://broker-two:4222" + current.Store(&u) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StateCommit)) + ready("a", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + Expect(row().State).To(Equal(cluster.StateStable)) + Expect(row().NATSURL).To(Equal("nats://broker-one:4222")) + }) + + It("never copies a credential or a query string into the row", func() { + u := "nats://user:secret@broker-one:4222?token=abc" + current.Store(&u) + replica("a", "", "") + got, _, err := request(cluster.CarrierNATS, false) + Expect(err).ToNot(HaveOccurred()) + Expect(got.NATSURL).To(Equal("nats://broker-one:4222")) + }) + }) + It("changes nothing and returns the report when the preflight blocks", func() { replica("a", "", "") workers.set(cluster.WorkerInfo{ID: "old", Name: "old"}) From 289c82dd086d4197076c90d98e2fb58b534c1cb2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:10:57 +0000 Subject: [PATCH 135/162] fix(jobs): stop a handler whose claim is gone, free the claims of a restart, end a poison row Three defects of the claim queue let work run twice, never run, or run for ever. - A replica whose instances row went missing kept running its handler while a peer reaped the claim and ran the job again. The consumer now checks every tick (once for all kinds) that the replica is live and lists the claims it holds. A handler whose claim is gone gets its context cancelled and does not settle the row. When the replica registers again before a peer reaped the row, it gives the row back, because the reap never frees the rows of a live id. - A process that restarted with the same instance id never ran the claims of the earlier process: the id is live again at once, so the reap skipped them. A loop now releases the rows of its own id when it starts. - A row that failed every time was released every minute for ever. Failures that say something about the work are now counted. After MaxFailures (default 10) the consumer records the job as failed through OnExhausted and deletes the row. A release for lack of a worker, a route or a free slot is not counted, so a fleet that is down keeps its queue. The reap, the liveness check and the purge of old migrated rows run once for each interval and not once for each kind. The delivery contract is at least once, and the doc comment of the consumer says so. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/jobs/claim.go | 127 ++++++++++- core/services/jobs/claim_queue.go | 244 +++++++++++++++++++--- core/services/jobs/claim_recovery_test.go | 215 +++++++++++++++++++ core/services/jobs/claim_test.go | 21 ++ 4 files changed, 577 insertions(+), 30 deletions(-) create mode 100644 core/services/jobs/claim_recovery_test.go diff --git a/core/services/jobs/claim.go b/core/services/jobs/claim.go index 870dd0eaf857..d45eb21c949e 100644 --- a/core/services/jobs/claim.go +++ b/core/services/jobs/claim.go @@ -53,6 +53,12 @@ type WorkClaim struct { // that no longer matches the row, so it cannot settle the claim of the // replica that took the work over. Attempts int + // Failures counts the releases that followed a failure which says something + // about the work itself: a worker was reached and the run failed on it. A + // release for lack of a worker, a route or a free slot does not count. The + // count is what ends a row that can never be run (see + // ClaimConsumerConfig.MaxFailures). + Failures int // NotBefore is the earliest the row may be claimed again. The database clock // stamps it. NULL means now. NotBefore *time.Time `gorm:"index"` @@ -134,6 +140,7 @@ func EnqueueClaim(ctx context.Context, db *gorm.DB, kind messaging.WorkKind, pay "payload": raw, "claimed_by": "", "attempts": 0, + "failures": 0, "created_at": gorm.Expr("now()"), }).Error; err != nil { return "", fmt.Errorf("enqueueing a %s claim: %w", kind, err) @@ -161,7 +168,7 @@ WHERE id = ( FOR UPDATE SKIP LOCKED LIMIT 1 ) -RETURNING id, kind, state, payload, claimed_by, claimed_at, attempts, not_before, created_at` +RETURNING id, kind, state, payload, claimed_by, claimed_at, attempts, failures, not_before, created_at` // ClaimNext takes the oldest pending row of kind and marks it as held by owner. // It returns ErrNoWork when there is none. @@ -239,6 +246,71 @@ func ReleaseClaim(ctx context.Context, db *gorm.DB, claim *WorkClaim) (bool, err return res.RowsAffected == 1, nil } +// releaseClaimFailedSQL is releaseClaimSQL for a failure that says something +// about the work, which it counts. +const releaseClaimFailedSQL = `UPDATE ` + claimsTable + ` +SET state = '` + ClaimPending + `', claimed_by = '', claimed_at = NULL, attempts = attempts + 1, failures = failures + 1, + not_before = now() + make_interval(secs => LEAST(?, ? * power(2, LEAST(attempts, ?)))) +WHERE id = ? AND state = '` + ClaimClaimed + `' AND claimed_by = ? AND attempts = ?` + +// ReleaseClaimFailed is ReleaseClaim for a run that reached a worker and failed +// there. It counts the failure, so that a row which fails every time can be ended. +func ReleaseClaimFailed(ctx context.Context, db *gorm.DB, claim *WorkClaim) (bool, error) { + if err := requirePostgres(db, "releasing a claim"); err != nil { + return false, err + } + res := db.WithContext(ctx).Exec(releaseClaimFailedSQL, + claimBackoffCap.Seconds(), claimBackoffBase.Seconds(), claimBackoffMaxShift, + claim.ID, claim.ClaimedBy, claim.Attempts) + if res.Error != nil { + return false, fmt.Errorf("releasing claim %q: %w", claim.ID, res.Error) + } + return res.RowsAffected == 1, nil +} + +// releaseOwnedSQL returns the rows of a kind that a replica holds to the pool, with +// no wait. +const releaseOwnedSQL = `UPDATE ` + claimsTable + ` +SET state = '` + ClaimPending + `', claimed_by = '', claimed_at = NULL, attempts = attempts + 1, not_before = NULL +WHERE state = '` + ClaimClaimed + `' AND claimed_by = ? AND kind = ?` + +// ReleaseOwned returns every claim of kind that owner holds to the pool. A +// process that starts to claim calls it first: a row that its own id holds was +// held by an earlier process with that id, which is gone. The reap cannot see +// that, because the id is live again, and the row would be held for ever. +func ReleaseOwned(ctx context.Context, db *gorm.DB, owner string, kind messaging.WorkKind) (int64, error) { + if err := requirePostgres(db, "releasing the claims of this replica"); err != nil { + return 0, err + } + res := db.WithContext(ctx).Exec(releaseOwnedSQL, owner, string(kind)) + if res.Error != nil { + return 0, fmt.Errorf("releasing the claims held by %q: %w", owner, res.Error) + } + return res.RowsAffected, nil +} + +// HeldClaim names a claim that a replica holds now. +type HeldClaim struct { + ID string + Attempts int +} + +// HeldBy lists the claims that owner holds. A run in flight whose claim is not in +// the list has lost it: a reap gave the work to another replica. +func HeldBy(ctx context.Context, db *gorm.DB, owner string) ([]HeldClaim, error) { + if err := requirePostgres(db, "listing the claims of this replica"); err != nil { + return nil, err + } + var held []HeldClaim + if err := db.WithContext(ctx).Model(&WorkClaim{}). + Select("id, attempts"). + Where("state = ? AND claimed_by = ?", ClaimClaimed, owner). + Scan(&held).Error; err != nil { + return nil, fmt.Errorf("listing the claims held by %q: %w", owner, err) + } + return held, nil +} + // CompleteClaim deletes the row. It is what an answer of the worker does, // success or failure: the work ran and said what happened, so it must not run // again. It reports false when the claim was no longer the holder's. @@ -309,19 +381,23 @@ func OwnerIsLive(ctx context.Context, db *gorm.DB, owner string, liveness time.D // migratePendingSQL takes the pending rows out of the queue. Rows that a replica // holds are not touched: the replica drives them to their end where they are. +// +// claimed_at holds the time of the sweep, for the purge of the rows that stay +// behind. const migratePendingSQL = `UPDATE ` + claimsTable + ` -SET state = '` + ClaimMigrated + `' +SET state = '` + ClaimMigrated + `', claimed_at = now() WHERE state = '` + ClaimPending + `' -RETURNING id, kind, state, payload, claimed_by, claimed_at, attempts, not_before, created_at` +RETURNING id, kind, state, payload, claimed_by, claimed_at, attempts, failures, not_before, created_at` // MigratePending takes every pending row out of the queue and returns it, for a // change to a carrier that has its own queue. A row is migrated once: the // statement is a compare-and-set on the state, so it never goes to a claimant and // to the sweep, and a second sweep returns nothing. // -// The caller publishes the rows after this returns. A sweeper that dies between -// the two loses the rows, as the queue of the new carrier loses them, and the -// alternative, to publish first, would run a job twice. +// The caller publishes the rows after this returns, and puts back the ones it +// could not publish (RestorePending). A sweeper that dies between the two loses +// the rows, as the queue of the new carrier loses them, and the alternative, to +// publish first, would run a job twice. func MigratePending(ctx context.Context, db *gorm.DB) ([]*WorkClaim, error) { if err := requirePostgres(db, "migrating pending work"); err != nil { return nil, err @@ -332,3 +408,42 @@ func MigratePending(ctx context.Context, db *gorm.DB) ([]*WorkClaim, error) { } return rows, nil } + +// RestorePending puts migrated rows back in the queue. It is for the rows that a +// change of carrier took out and could not hand to the new carrier: they wait in +// the table, and the next sweep takes them again, and none is lost. +func RestorePending(ctx context.Context, db *gorm.DB, ids []string) error { + if len(ids) == 0 { + return nil + } + if err := requirePostgres(db, "restoring migrated work"); err != nil { + return err + } + res := db.WithContext(ctx).Model(&WorkClaim{}). + Where("id IN ? AND state = ?", ids, ClaimMigrated). + Updates(map[string]any{"state": ClaimPending, "claimed_at": nil}) + if res.Error != nil { + return fmt.Errorf("restoring migrated work: %w", res.Error) + } + return nil +} + +// MigratedRetention is how long a row that a change of carrier moved stays in the +// table. It is only evidence of what moved. The payload of a row is the whole +// request of the work, so it does not stay for ever. +const MigratedRetention = 24 * time.Hour + +// PurgeMigrated deletes the rows that were migrated more than olderThan ago, with +// their payloads, and says how many. +func PurgeMigrated(ctx context.Context, db *gorm.DB, olderThan time.Duration) (int64, error) { + if err := requirePostgres(db, "purging migrated work"); err != nil { + return 0, err + } + res := db.WithContext(ctx).Exec(`DELETE FROM `+claimsTable+ + ` WHERE state = ? AND COALESCE(claimed_at, created_at) < now() - make_interval(secs => ?)`, + ClaimMigrated, olderThan.Seconds()) + if res.Error != nil { + return 0, fmt.Errorf("purging migrated work: %w", res.Error) + } + return res.RowsAffected, nil +} diff --git a/core/services/jobs/claim_queue.go b/core/services/jobs/claim_queue.go index 95e871ea1791..0b486d14557d 100644 --- a/core/services/jobs/claim_queue.go +++ b/core/services/jobs/claim_queue.go @@ -83,8 +83,39 @@ type ClaimConsumerConfig struct { // Liveness is the window of replica liveness. cluster.InstanceLiveness when // zero: it is the one definition of absence that the cluster has. Liveness time.Duration + // MaxFailures is how many failures that say something about the work a row + // may have before the consumer ends it (OnExhausted, then delete). + // DefaultClaimMaxFailures when zero. A release for lack of a worker, a route + // or a free slot is not a failure of the work and is not counted. + MaxFailures int + // OnExhausted is called for a row that has reached MaxFailures, to record in + // the job that it failed. The row is deleted when it returns nil, and kept + // (to be tried again) when it returns an error. It may be nil. + OnExhausted func(ctx context.Context, kind messaging.WorkKind, payload []byte, cause error) error } +// DefaultClaimMaxFailures is how many failures end a row. With the backoff of +// the releases this is about five minutes of trying. +const DefaultClaimMaxFailures = 10 + +var ( + // errClaimLost is the cause of the context of a handler whose claim another + // replica took over: the instances row of this replica was missing for long + // enough for the reap, and the work was given away. + errClaimLost = errors.New("jobs: the claim was taken over by another replica") + // errOwnerNotLive is the cause of the context of a handler when this replica + // is no longer live in the instances table. Its peers may reap its claims at + // any moment, so it stops the work instead of running it twice. + errOwnerNotLive = errors.New("jobs: this replica is no longer live in the cluster") + // errNoAnswerYet marks the error of a handler that says nothing about the + // work: no worker, no route, no free slot. The row goes back to the pool and + // the failure is not counted. + errNoAnswerYet = errors.New("jobs: the work could not be offered to a worker") +) + +// noAnswerYet marks err as one that says nothing about the work. +func noAnswerYet(err error) error { return fmt.Errorf("%w: %w", errNoAnswerYet, err) } + // DefaultClaimInterval is how often a consumer looks for work when no hint woke // it. const DefaultClaimInterval = 2 * time.Second @@ -97,12 +128,39 @@ const DefaultClaimInterval = 2 * time.Second // The outcome of a handler settles the row. A nil return completes it: the work // has an answer. An error releases it with a wait: nothing was learned about the // work, so it must be tried again, possibly by another replica. A panic is an -// error. An error that wraps ErrKeepClaim leaves the row alone. +// error. An error that wraps ErrKeepClaim leaves the row alone. A row whose +// failures reach MaxFailures is ended, and not released again. +// +// The contract is at least once. A handler can be called again for the same work: +// after the replica that held it died, after a release, and after the instances +// row of its replica was missing for as long as a peer needed to reap its claims +// (a stall of the process or of its link to the database for about the liveness +// window). In that last case the replica stops its handler, by cancelling the +// context that the handler was given, as soon as it sees that the claim is no +// longer its own; the work that the handler already did stays done. A handler for +// work that is not safe to repeat must record its progress and check it when it +// is called. // -// A handler can be called again for the same work: after the replica that held -// it died, after a release. It must tolerate a repeat. +// The loops of the kinds share what they ask the database every tick: whether +// this replica is live, which claims it holds, and the reap of the claims of +// departed replicas. Three loops ask once and not three times. type ClaimConsumer struct { cfg ClaimConsumerConfig + + checkMu sync.Mutex + checkAt time.Time + live bool + reapAt time.Time + + flightMu sync.Mutex + flights map[*flight]struct{} +} + +// flight is a run in progress: the claim it drives and how to stop it. +type flight struct { + id string + attempts int + cancel context.CancelCauseFunc } var _ messaging.WorkConsumer = (*ClaimConsumer)(nil) @@ -124,7 +182,10 @@ func NewClaimConsumer(cfg ClaimConsumerConfig) (*ClaimConsumer, error) { if cfg.Liveness <= 0 { cfg.Liveness = cluster.InstanceLiveness } - return &ClaimConsumer{cfg: cfg}, nil + if cfg.MaxFailures <= 0 { + cfg.MaxFailures = DefaultClaimMaxFailures + } + return &ClaimConsumer{cfg: cfg, flights: map[*flight]struct{}{}}, nil } // Consume claims the rows of kind and runs h on each, at most maxInFlight at a @@ -137,6 +198,7 @@ func (c *ClaimConsumer) Consume(ctx context.Context, kind messaging.WorkKind, ma } loopCtx, stop := context.WithCancel(ctx) l := &claimLoop{ + c: c, cfg: c.cfg, kind: kind, handler: h, @@ -162,6 +224,7 @@ func (c *ClaimConsumer) Consume(ctx context.Context, kind messaging.WorkKind, ma // claimLoop is one subscription: the loop that claims the rows of one kind. type claimLoop struct { + c *ClaimConsumer cfg ClaimConsumerConfig kind messaging.WorkKind handler messaging.WorkHandler @@ -197,6 +260,7 @@ func (l *claimLoop) hear(raw []byte) { func (l *claimLoop) run(ctx context.Context) { defer close(l.done) + l.releaseOwn(ctx) ticker := time.NewTicker(l.cfg.Interval) defer ticker.Stop() l.fill(ctx, true) @@ -231,6 +295,112 @@ func (l *claimLoop) free() { } } +// releaseOwn returns the claims of this kind that the id of this replica holds, +// when the loop starts. This process has driven none of them yet: they belong to +// an earlier process that had the same id (LOCALAI_INSTANCE_ID is stable across +// restarts), and the reap would never free them, because the id is live. +func (l *claimLoop) releaseOwn(ctx context.Context) { + n, err := ReleaseOwned(ctx, l.cfg.DB, l.cfg.Owner, l.kind) + switch { + case err != nil: + if ctx.Err() == nil { + xlog.Warn("Could not release the claims that an earlier process of this replica left", "kind", l.kind, "error", err) + } + case n > 0: + xlog.Info("Released work that an earlier process of this replica was driving", "kind", l.kind, "count", n) + } +} + +// register adds a run in progress. +func (c *ClaimConsumer) register(claim *WorkClaim, cancel context.CancelCauseFunc) *flight { + f := &flight{id: claim.ID, attempts: claim.Attempts, cancel: cancel} + c.flightMu.Lock() + c.flights[f] = struct{}{} + c.flightMu.Unlock() + return f +} + +func (c *ClaimConsumer) unregister(f *flight) { + c.flightMu.Lock() + delete(c.flights, f) + c.flightMu.Unlock() +} + +func (c *ClaimConsumer) snapshot() []*flight { + c.flightMu.Lock() + defer c.flightMu.Unlock() + out := make([]*flight, 0, len(c.flights)) + for f := range c.flights { + out = append(out, f) + } + return out +} + +// tick does what every loop needs once for each interval, and lets the loops that +// ask within half an interval reuse the answer: it checks that this replica is +// live, stops the runs that lost their claim, and, on a poll, frees the claims of +// departed replicas and purges old migrated rows. It says whether this replica +// may claim. +func (c *ClaimConsumer) tick(ctx context.Context, poll bool) (bool, error) { + c.checkMu.Lock() + defer c.checkMu.Unlock() + fresh := !c.checkAt.IsZero() && time.Since(c.checkAt) < c.cfg.Interval/2 + if !fresh { + // The runs are counted before the question is asked. A claim that is taken + // after it is not judged by it. + runs := c.snapshot() + live, err := OwnerIsLive(ctx, c.cfg.DB, c.cfg.Owner, c.cfg.Liveness) + if err != nil { + return false, err + } + c.live, c.checkAt = live, time.Now() + if !live { + // Its peers may reap the claims of this replica at any moment, and then + // the work runs on two replicas. Stop it here. + for _, f := range runs { + f.cancel(errOwnerNotLive) + } + } else if len(runs) > 0 { + held, err := HeldBy(ctx, c.cfg.DB, c.cfg.Owner) + if err != nil { + if ctx.Err() == nil { + xlog.Warn("Could not check which claims this replica still holds", "error", err) + } + } else { + still := make(map[string]int, len(held)) + for _, h := range held { + still[h.ID] = h.Attempts + } + for _, f := range runs { + if attempts, ok := still[f.id]; !ok || attempts != f.attempts { + f.cancel(errClaimLost) + } + } + } + } + } + if poll && c.live && time.Since(c.reapAt) >= c.cfg.Interval/2 { + c.reapAt = time.Now() + if released, err := ReapAbandoned(ctx, c.cfg.DB, c.cfg.Liveness); err != nil { + // Reported and not returned: a reap that could not run says nothing + // about the claim this tick is about to take. + if ctx.Err() == nil { + xlog.Warn("Could not release the claims of departed replicas", "error", err) + } + } else if released > 0 { + xlog.Info("Released work claimed by replicas that are no longer live", "count", released) + } + if purged, err := PurgeMigrated(ctx, c.cfg.DB, MigratedRetention); err != nil { + if ctx.Err() == nil { + xlog.Warn("Could not purge the work that a change of carrier moved", "error", err) + } + } else if purged > 0 { + xlog.Info("Purged work that a change of carrier moved long ago", "count", purged) + } + } + return c.live, nil +} + // fill claims rows into the free slots. On a poll it first looks after the // claims that departed replicas left. func (l *claimLoop) fill(ctx context.Context, poll bool) { @@ -239,7 +409,7 @@ func (l *claimLoop) fill(ctx context.Context, poll bool) { // is a degraded state, but not one for holding a claim: the reap would take // the claim out from under it while the work still ran, and the work would // run twice. Refusing to claim heals by itself when membership registers. - live, err := OwnerIsLive(ctx, l.cfg.DB, l.cfg.Owner, l.cfg.Liveness) + live, err := l.c.tick(ctx, poll) if err != nil { if ctx.Err() == nil { xlog.Warn("Could not check whether this replica may claim work", "kind", l.kind, "error", err) @@ -259,18 +429,6 @@ func (l *claimLoop) fill(ctx context.Context, poll bool) { xlog.Info("This replica is live in the cluster again and claims queued work", "instance", l.cfg.Owner, "kind", l.kind) } - if poll { - if released, err := ReapAbandoned(ctx, l.cfg.DB, l.cfg.Liveness); err != nil { - // Reported and not returned: a reap that could not run says nothing - // about the claim this tick is about to take. - if ctx.Err() == nil { - xlog.Warn("Could not release the claims of departed replicas", "error", err) - } - } else if released > 0 { - xlog.Info("Released work claimed by replicas that are no longer live", "count", released) - } - } - for ctx.Err() == nil { if !l.take() { return // this replica drives as much as it may @@ -292,22 +450,60 @@ func (l *claimLoop) fill(ctx context.Context, poll bool) { func (l *claimLoop) drive(claim *WorkClaim) { defer l.wg.Done() defer l.free() - err := l.call(claim) + hctx, cancel := context.WithCancelCause(l.parent) + f := l.c.register(claim, cancel) + defer l.c.unregister(f) + defer cancel(nil) + err := l.call(hctx, claim) + cause := context.Cause(hctx) // The settle must outlive the stop of the replica. - ctx, cancel := context.WithTimeout(context.WithoutCancel(l.parent), settleTimeout) - defer cancel() + ctx, cancelSettle := context.WithTimeout(context.WithoutCancel(l.parent), settleTimeout) + defer cancelSettle() switch { + case err != nil && errors.Is(cause, errClaimLost): + // The row belongs to another replica now. Settling it would be wrong, and + // the fence of the attempt count would refuse it anyway. + xlog.Warn("Stopped a run whose claim another replica took over; the work runs there", + "claim", claim.ID, "kind", claim.Kind) + case err != nil && errors.Is(cause, errOwnerNotLive): + // This replica came back, or may come back, before a peer reaped the row. + // A row that stays claimed by a live id is never reaped, so give it back. + xlog.Warn("Stopped a run: this replica was not live in the cluster. Returning the claim to the queue", + "claim", claim.ID, "kind", claim.Kind) + l.settle(ReleaseClaim(ctx, l.cfg.DB, claim)) case err == nil: l.settle(CompleteClaim(ctx, l.cfg.DB, claim)) case errors.Is(err, ErrKeepClaim): xlog.Warn("Leaving a claim as it is: the work ran and its answer could not be recorded", "claim", claim.ID, "kind", claim.Kind, "error", err) - default: - xlog.Warn("Releasing a claim whose work obtained no answer; it will be tried again after a wait", + case errors.Is(err, errNoAnswerYet): + xlog.Warn("Releasing a claim whose work could not be offered to a worker; it will be tried again after a wait", "claim", claim.ID, "kind", claim.Kind, "attempt", claim.Attempts+1, "error", err) l.settle(ReleaseClaim(ctx, l.cfg.DB, claim)) + case claim.Failures+1 >= l.cfg.MaxFailures: + l.exhaust(ctx, claim, err) + default: + xlog.Warn("Releasing a claim whose work obtained no answer; it will be tried again after a wait", + "claim", claim.ID, "kind", claim.Kind, "attempt", claim.Attempts+1, "failure", claim.Failures+1, "of", l.cfg.MaxFailures, "error", err) + l.settle(ReleaseClaimFailed(ctx, l.cfg.DB, claim)) + } +} + +// exhaust ends a row that failed MaxFailures times. The job is recorded as +// failed first and the row is deleted after, so a row is never deleted with its +// job still running. +func (l *claimLoop) exhaust(ctx context.Context, claim *WorkClaim, cause error) { + xlog.Error("Giving up on queued work: it failed too many times", + "claim", claim.ID, "kind", claim.Kind, "failures", claim.Failures+1, "error", cause) + if l.cfg.OnExhausted != nil { + if err := l.cfg.OnExhausted(ctx, l.kind, claim.Payload, cause); err != nil { + xlog.Error("Could not record that queued work failed; it will be tried again", "claim", claim.ID, "error", err) + l.settle(ReleaseClaimFailed(ctx, l.cfg.DB, claim)) + return + } } + l.settle(CompleteClaim(ctx, l.cfg.DB, claim)) } func (l *claimLoop) settle(taken bool, err error) { @@ -323,13 +519,13 @@ func (l *claimLoop) settle(taken bool, err error) { // call runs the handler and turns a panic into an error, so that one bad payload // cannot take the replica down with it. -func (l *claimLoop) call(claim *WorkClaim) (err error) { +func (l *claimLoop) call(ctx context.Context, claim *WorkClaim) (err error) { defer func() { if p := recover(); p != nil { err = fmt.Errorf("the handler panicked: %v", p) } }() - return l.handler(l.parent, claim.Payload, l.cfg.Events) + return l.handler(ctx, claim.Payload, l.cfg.Events) } // Unsubscribe stops claiming and waits for the handlers that run. diff --git a/core/services/jobs/claim_recovery_test.go b/core/services/jobs/claim_recovery_test.go new file mode 100644 index 000000000000..0c547b9899f3 --- /dev/null +++ b/core/services/jobs/claim_recovery_test.go @@ -0,0 +1,215 @@ +package jobs + +import ( + "context" + "errors" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" +) + +var _ = Describe("The claim queue when a replica loses its claims or its process", func() { + var ( + ctx context.Context + db *gorm.DB + bus *testutil.FakeBus + ) + + BeforeEach(func() { + c, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + ctx = c + db = newClaimDB() + bus = testutil.NewFakeBus() + }) + + consume := func(c *ClaimConsumer, kind messaging.WorkKind, h messaging.WorkHandler) { + GinkgoHelper() + sub, err := c.Consume(ctx, kind, 0, h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + } + // skipWaits clears the wait that a release stamps, so that a spec does not wait + // for the backoff of the queue. + skipWaits := func() { + stop := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + for { + select { + case <-stop: + return + case <-time.After(20 * time.Millisecond): + _ = db.Exec(`UPDATE work_claims SET not_before = NULL`).Error + } + } + }() + DeferCleanup(func() { close(stop); <-done }) + } + rows := func() int64 { + var n int64 + Expect(db.Model(&WorkClaim{}).Count(&n).Error).To(Succeed()) + return n + } + + Describe("a replica whose instances row went missing while it ran a job", func() { + It("stops the handler of the job that a peer took over", func() { + started := make(chan struct{}) + ended := make(chan error, 1) + var ranOnB atomic.Int32 + + a := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(a, messaging.WorkMCPCI, func(hctx context.Context, _ []byte, _ messaging.Publisher) error { + close(started) + <-hctx.Done() + ended <- context.Cause(hctx) + return hctx.Err() + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "work")).To(Succeed()) + Eventually(started, 5*time.Second).Should(BeClosed()) + + // The row of A is missing for the liveness window, and B reaps its claim. + kill(db, "replica-a") + b := newConsumer(db, consumerOptions{owner: "replica-b", hints: bus, interval: 100 * time.Millisecond}) + consume(b, messaging.WorkMCPCI, func(context.Context, []byte, messaging.Publisher) error { + ranOnB.Add(1) + return nil + }) + Eventually(ranOnB.Load, 10*time.Second).Should(BeEquivalentTo(1)) + + var cause error + Eventually(ended, 5*time.Second).Should(Receive(&cause), "the handler on A is told to stop, it does not run on") + Expect(errors.Is(cause, errClaimLost) || errors.Is(cause, errOwnerNotLive)).To(BeTrue(), "cause: %v", cause) + Eventually(rows, 5*time.Second).Should(BeZero()) + }) + + It("gives the claim back when the replica registers again before a peer reaped it, and the job runs again", func() { + firstEnded := make(chan struct{}) + var calls atomic.Int32 + a := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(a, messaging.WorkMCPCI, func(hctx context.Context, _ []byte, _ messaging.Publisher) error { + if calls.Add(1) == 1 { + <-hctx.Done() + close(firstEnded) + return hctx.Err() + } + return nil + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "work")).To(Succeed()) + Eventually(calls.Load, 5*time.Second).Should(BeEquivalentTo(1)) + + kill(db, "replica-a") + Eventually(firstEnded, 5*time.Second).Should(BeClosed()) + live(db, "replica-a") // its heartbeat writes the row again + + // Without the release the row stays claimed by a live id for ever: the reap + // only frees the claims of ids that are gone. + Eventually(calls.Load, 15*time.Second).Should(BeEquivalentTo(2)) + Eventually(rows, 5*time.Second).Should(BeZero()) + }) + }) + + Describe("a replica that restarts with the same instance id", func() { + It("recovers the job that the earlier process held, inside the liveness window", func() { + started := make(chan struct{}) + hold := make(chan struct{}) + var ranAgain atomic.Int32 + + // The first process claims the job and is gone: nothing settles the row. + first := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + sub, err := first.Consume(ctx, messaging.WorkMCPCI, 0, func(context.Context, []byte, messaging.Publisher) error { + close(started) + <-hold + return ErrKeepClaim // as a process that died would leave it + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + DeferCleanup(func() { close(hold) }) // runs before the Unsubscribe above, which waits for the handler + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "work")).To(Succeed()) + Eventually(started, 5*time.Second).Should(BeClosed()) + + // The new process has the same id, and the id never left the instances table. + second := newConsumer(db, consumerOptions{owner: "replica-a", hints: bus, interval: 100 * time.Millisecond}) + consume(second, messaging.WorkMCPCI, func(context.Context, []byte, messaging.Publisher) error { + ranAgain.Add(1) + return nil + }) + Eventually(ranAgain.Load, 10*time.Second).Should(BeEquivalentTo(1)) + Eventually(rows, 5*time.Second).Should(BeZero()) + }) + }) + + Describe("a row that can never be run", func() { + It("ends after the failures that the consumer allows, records the failure and deletes the row", func() { + skipWaits() + var calls atomic.Int32 + var gaveUp atomic.Pointer[string] + live(db, "replica-a") + c, err := NewClaimConsumer(ClaimConsumerConfig{ + DB: db, Owner: "replica-a", Events: bus, Hints: bus, Interval: 50 * time.Millisecond, + MaxFailures: 3, + OnExhausted: func(_ context.Context, _ messaging.WorkKind, payload []byte, cause error) error { + msg := string(payload) + ": " + cause.Error() + gaveUp.Store(&msg) + return nil + }, + }) + Expect(err).ToNot(HaveOccurred()) + consume(c, messaging.WorkMCPCI, func(context.Context, []byte, messaging.Publisher) error { + calls.Add(1) + return errors.New("the worker refused the payload") + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "poison")).To(Succeed()) + + Eventually(rows, 15*time.Second).Should(BeZero()) + Expect(calls.Load()).To(BeEquivalentTo(3), "it ran as often as the consumer allows and not again") + Expect(*gaveUp.Load()).To(ContainSubstring("the worker refused the payload")) + }) + + It("keeps the row when the failure cannot be recorded, and tries again", func() { + skipWaits() + var recorded atomic.Int32 + live(db, "replica-a") + c, err := NewClaimConsumer(ClaimConsumerConfig{ + DB: db, Owner: "replica-a", Events: bus, Interval: 50 * time.Millisecond, MaxFailures: 1, + OnExhausted: func(context.Context, messaging.WorkKind, []byte, error) error { + if recorded.Add(1) == 1 { + return errors.New("the database is away") + } + return nil + }, + }) + Expect(err).ToNot(HaveOccurred()) + consume(c, messaging.WorkMCPCI, func(context.Context, []byte, messaging.Publisher) error { + return errors.New("failed") + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "poison")).To(Succeed()) + Eventually(rows, 15*time.Second).Should(BeZero()) + Expect(recorded.Load()).To(BeNumerically(">=", 2)) + }) + + It("does not count a release for lack of a worker as a failure", func() { + skipWaits() + var calls atomic.Int32 + live(db, "replica-a") + c, err := NewClaimConsumer(ClaimConsumerConfig{ + DB: db, Owner: "replica-a", Events: bus, Interval: 50 * time.Millisecond, MaxFailures: 2, + }) + Expect(err).ToNot(HaveOccurred()) + consume(c, messaging.WorkMCPCI, func(context.Context, []byte, messaging.Publisher) error { + calls.Add(1) + return noAnswerYet(errors.New("no agent worker is connected")) + }) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, "waits")).To(Succeed()) + Eventually(calls.Load, 15*time.Second).Should(BeNumerically(">=", 5)) + Expect(rows()).To(BeEquivalentTo(1), "a fleet that is down keeps its queued work") + }) + }) +}) diff --git a/core/services/jobs/claim_test.go b/core/services/jobs/claim_test.go index b315af257f2e..f7af75ea4e88 100644 --- a/core/services/jobs/claim_test.go +++ b/core/services/jobs/claim_test.go @@ -200,6 +200,27 @@ var _ = Describe("The claim table", func() { }) }) + Describe("the poll query", func() { + It("reads the rows of a kind through the index of the pick, in the order of the queue", func() { + db := newClaimDB() + for i := range 50 { + _, err := EnqueueClaim(context.Background(), db, messaging.WorkMCPCI, i) + Expect(err).ToNot(HaveOccurred()) + } + Expect(db.Exec("ANALYZE " + claimsTable).Error).To(Succeed()) + var plan []string + Expect(db.Transaction(func(tx *gorm.DB) error { + if err := tx.Exec("SET LOCAL enable_seqscan = off").Error; err != nil { + return err + } + return tx.Raw(`EXPLAIN SELECT id FROM ` + claimsTable + ` + WHERE state = 'pending' AND kind = 'mcp-ci' AND (not_before IS NULL OR not_before <= now()) + ORDER BY created_at, id LIMIT 1`).Scan(&plan).Error + })).To(Succeed()) + Expect(strings.Join(plan, "\n")).To(ContainSubstring("idx_work_claims_pick")) + }) + }) + Describe("CompleteClaim", func() { It("deletes the row", func() { id, _ := EnqueueClaim(ctx, db, messaging.WorkAgentRun, 1) From 13d48c1a1b749d3c82c172eb3e0dbfcd74b4bbcb Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:10:57 +0000 Subject: [PATCH 136/162] fix(jobs): do not run a job again once it started, and honour a cancel that came early The driver released the claim of any run whose stream broke, so a job that had already started a pipeline or an agent run was started a second time. It also dropped a cancel that came before the claim, and ran again a job that was cancelled while it waited for a free worker. - Once a line of the run has come back, a broken link closes the job as failed and completes the claim. A claim is released only when the run did not start: no worker, no route, no free slot, an old worker, or a carrier that was released first. Before, the end of a drain completed the claim of a run that never started and dropped the work. - A worker with no route to it, or one that answers 404 for the verb, is skipped like a busy one: the next worker gets the run with no wait. - Dispatcher.Cancel marks the job as cancelled before it sends the event. The driver reads the status before it dispatches and before each worker it offers the run to. A cancel that ends the context while a worker is picked is a cancel and not a missing worker. - The sentinels for "no route" and "verb not served" now live in workerctl, and the nodes package keeps its names as aliases, so that the jobs package can read them without a new import. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/jobs/claim_driver.go | 136 ++++++++++++++++---- core/services/jobs/claim_driver_test.go | 151 ++++++++++++++++++++++- core/services/jobs/dispatch_loop.go | 4 + core/services/jobs/dispatch_loop_test.go | 46 ++++++- core/services/jobs/dispatcher.go | 10 +- core/services/jobs/dispatcher_test.go | 36 ++++++ core/services/jobs/store.go | 9 ++ core/services/nodes/control_nats.go | 2 +- core/services/nodes/unloader.go | 2 +- core/services/workerctl/request.go | 10 ++ 10 files changed, 372 insertions(+), 34 deletions(-) diff --git a/core/services/jobs/claim_driver.go b/core/services/jobs/claim_driver.go index 3bbee9ee9dfb..7044d68405ce 100644 --- a/core/services/jobs/claim_driver.go +++ b/core/services/jobs/claim_driver.go @@ -168,10 +168,25 @@ func (d *AgentDriver) handler(kind messaging.WorkKind, verb string) messaging.Wo } } -// drive hands one claim to a worker and returns what settles it: nil when the -// worker answered (or the job was cancelled), ErrKeepClaim when the worker -// answered and the answer could not be written, and any other error when nothing -// was learned about the work. +// JobStatusReader reads the status of a job. *JobStore is one. The driver reads +// it before it hands a job to a worker, so that a job that was cancelled while it +// waited in the queue is not run. +type JobStatusReader interface { + JobStatus(jobID string) (string, error) +} + +// drive hands one claim to a worker and returns what settles it: +// +// - nil when the worker answered, when the job was cancelled, and when a run +// that had started lost its link (the job is then closed as failed); +// - ErrKeepClaim when the worker answered and the answer could not be written; +// - an error marked noAnswerYet when the run was not started: no worker, no +// route, no free slot, an old worker, or the carrier was released first; +// - any other error when a worker was reached and the call failed before a +// single line of the run came back. +// +// A run that has started is never released. Releasing it would start the job a +// second time, and a job may run a CI pipeline or call an agent with side effects. func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb string, payload []byte) error { jobID := "" if kind == messaging.WorkMCPCI { @@ -181,49 +196,122 @@ func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb s var cancel context.CancelFunc ctx, cancel = context.WithCancel(ctx) defer cancel() + // Registered before the status is read: a cancel that comes after the + // read ends the run, and one that came before is in the status. d.cancels.Register(jobID, cancel) defer d.cancels.Deregister(jobID) defer d.cancelled.remove(jobID) } + // started becomes true when the first line of the run comes back. From then on + // the job may have done work, and the claim is not released. + var started atomic.Bool tried := make(map[string]bool, maxDrivePicks) var last error for range maxDrivePicks { + if d.cancelledNow(ctx, jobID) { + return d.persistCancelled(jobID) + } nodeID, nodeType, err := d.cfg.Picker.PickConnectedExcluding(ctx, tried) if err != nil { - if carrierReleased(ctx) { - return d.cutOff(jobID) + switch { + case carrierReleased(ctx): + return noAnswerYet(fmt.Errorf("the carrier was released before the %s claim started: %w", kind, err)) + case d.cancelledNow(ctx, jobID): + return d.persistCancelled(jobID) + case last != nil: + // Every worker that was offered the run had no slot, no route or an + // old build. That is the evidence, and the end of the list is not. + return noAnswerYet(last) } - if last != nil { - // Every worker that was offered the run had no slot. That is the - // evidence, and the end of the list is not. - return last - } - return fmt.Errorf("picking an agent worker for a %s claim: %w", kind, err) + return noAnswerYet(fmt.Errorf("picking an agent worker for a %s claim: %w", kind, err)) } tried[nodeID] = true var reply workerctl.RunReply err = d.cfg.Control.CallStreaming(ctx, nodeID, verb, json.RawMessage(payload), &reply, - func(subject string, raw json.RawMessage) { d.rebroadcast(nodeType, subject, raw) }) + func(subject string, raw json.RawMessage) { + started.Store(true) + d.rebroadcast(nodeType, subject, raw) + }) switch { case err == nil: return d.persistTerminal(&reply) case carrierReleased(ctx): - return d.cutOff(jobID) - case errors.Is(err, workerctl.ErrWorkerBusy): - last = fmt.Errorf("offering a %s claim to agent worker %q: %w", kind, nodeID, err) - continue + if started.Load() { + return d.cutOff(jobID) + } + return noAnswerYet(fmt.Errorf("the carrier was released before the %s claim started: %w", kind, err)) case jobID != "" && d.cancelled.has(jobID) && ctx.Err() != nil: // The stream was ended by the cancel of the job. The worker's run dies // with its request, and the job is closed here, because nothing else // will. - return d.persistTerminal(&workerctl.RunReply{JobID: jobID, Status: "cancelled", Error: "cancelled"}) + return d.persistCancelled(jobID) + case !started.Load() && (errors.Is(err, workerctl.ErrWorkerBusy) || + errors.Is(err, workerctl.ErrNoRoute) || errors.Is(err, workerctl.ErrVerbNotServed)): + // The run did not start on this worker: it has no slot, no route is + // held to it, or it is too old to serve the verb. The next worker may + // take it, and the claim waits for no backoff. + last = fmt.Errorf("offering a %s claim to agent worker %q: %w", kind, nodeID, err) + continue + case started.Load(): + return d.failStarted(jobID, nodeID, err) default: return fmt.Errorf("driving a %s claim on agent worker %q: %w", kind, nodeID, err) } } - return last + return noAnswerYet(last) +} + +// cancelledNow says whether the job was cancelled: this replica ended its run for +// that reason, or the job row says so. A job cancelled while it waited in the +// queue has no run to end, so the row is the only place that knows. +func (d *AgentDriver) cancelledNow(ctx context.Context, jobID string) bool { + if jobID == "" { + return false + } + if d.cancelled.has(jobID) && ctx.Err() != nil { + return true + } + reader, ok := d.cfg.Store.(JobStatusReader) + if !ok { + return false + } + status, err := reader.JobStatus(jobID) + return err == nil && status == "cancelled" +} + +// persistCancelled closes a job as cancelled. The claim is then completed: a +// cancel is an answer. +func (d *AgentDriver) persistCancelled(jobID string) error { + return d.persistTerminal(&workerctl.RunReply{JobID: jobID, Status: "cancelled", Error: "cancelled"}) +} + +// failStarted settles the claim of a run that had started when its link broke. +// The claim is completed and not released, because the run may have done work +// that a second run would repeat. The job is closed as failed, so that it does +// not stay running until the reaper finds it. A run with no job is a chat, whose +// output reached the user as events; there is nothing to close. +func (d *AgentDriver) failStarted(jobID, nodeID string, cause error) error { + xlog.Error("The link to an agent worker broke while a run was in progress; the run is not started again", + "job", jobID, "node", nodeID, "error", cause) + if jobID == "" { + return nil + } + return d.persistTerminal(&workerctl.RunReply{ + JobID: jobID, Status: "failed", + Error: "the link to the worker broke while the job ran: " + cause.Error(), + }) +} + +// failExhausted records that a job could not be run after the failures that the +// consumer allows. It is the OnExhausted of the consumer. +func (d *AgentDriver) failExhausted(_ context.Context, _ messaging.WorkKind, payload []byte, cause error) error { + jobID := jobIDOf(payload) + if jobID == "" || d.cfg.Store == nil { + return nil + } + return d.cfg.Store.UpdateJobStatus(jobID, "failed", "", "the job could not be run after repeated failures: "+cause.Error()) } // carrierReleased reports that ctx ended because the carrier of the loop was @@ -233,10 +321,12 @@ func carrierReleased(ctx context.Context) bool { return errors.Is(context.Cause(ctx), messaging.ErrCarrierReleased) } -// cutOff settles the claim of a run that the end of the drain cut off. The claim -// is completed, and not released: the run did start, so a carrier that took over -// the queue must not start it a second time. The job is left as it is, and the -// reaper fails it if it stays running, as it does for a job that lost its worker. +// cutOff settles the claim of a run that the end of the drain cut off after it had +// started. The claim is completed, and not released: the run did start, so a +// carrier that took over the queue must not start it a second time. The job is +// left as it is, and the reaper fails it if it stays running, as it does for a job +// that lost its worker. A run that had not started is released instead, and the +// carrier that takes over runs it. func (d *AgentDriver) cutOff(jobID string) error { xlog.Warn("A run was cut off by the end of the drain of its carrier; its job is left to the reaper", "job", jobID) return nil diff --git a/core/services/jobs/claim_driver_test.go b/core/services/jobs/claim_driver_test.go index b9364ddbd67a..f4a83327e1d3 100644 --- a/core/services/jobs/claim_driver_test.go +++ b/core/services/jobs/claim_driver_test.go @@ -95,9 +95,17 @@ func (b *fakeBroadcast) Handle(nodeType, subject string, raw json.RawMessage) bo type terminal struct{ jobID, status, result, errMsg string } type fakeStore struct { - mu sync.Mutex - got []terminal - fail error + mu sync.Mutex + got []terminal + fail error + statuses map[string]string +} + +// JobStatus makes the fake a JobStatusReader. +func (s *fakeStore) JobStatus(jobID string) (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.statuses[jobID], nil } func (s *fakeStore) UpdateJobStatus(jobID, status, result, errMsg string) error { @@ -274,6 +282,143 @@ var _ = Describe("The agent driver", func() { }) }) + Describe("a run whose link breaks", func() { + // startedThen plays a worker that sends a line of the run and then loses the link. + startedThen := func(err error) { + control.do = func(_ context.Context, _, _ string, onProgress func(string, json.RawMessage)) (workerctl.RunReply, error) { + onProgress("", json.RawMessage(`{"tick":1}`)) + return workerctl.RunReply{}, err + } + } + + It("does not release the claim of a job that had started, and closes the job as failed", func() { + startedThen(errors.New("the stream ended before its reply line")) + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed(), "a release would run the job a second time") + Expect(control.seen()).To(HaveLen(1)) + Expect(store.terminals()).To(HaveLen(1)) + Expect(store.terminals()[0].jobID).To(Equal("j1")) + Expect(store.terminals()[0].status).To(Equal("failed")) + Expect(store.terminals()[0].errMsg).To(ContainSubstring("link to the worker broke")) + }) + + It("keeps the claim when the failure of a started job cannot be written", func() { + startedThen(errors.New("unexpected EOF")) + store.fail = errors.New("database is away") + err := handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil) + Expect(errors.Is(err, ErrKeepClaim)).To(BeTrue()) + }) + + It("completes the claim of a chat run that had started, because it has no job to close", func() { + startedThen(errors.New("unexpected EOF")) + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).To(Succeed()) + Expect(store.terminals()).To(BeEmpty()) + }) + + It("still releases when the link breaks before any line came back", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, errors.New("connection reset") + } + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).ToNot(Succeed()) + Expect(store.terminals()).To(BeEmpty()) + }) + }) + + Describe("a worker that cannot take the run", func() { + for name, cause := range map[string]error{ + "has no route to it": fmt.Errorf("control request: %w", workerctl.ErrNoRoute), + "is too old to serve the verb": fmt.Errorf("control request: %w", workerctl.ErrVerbNotServed), + "has no free slot": workerctl.ErrWorkerBusy, + } { + It("moves on to the next worker when the first "+name, func() { + control.do = func(_ context.Context, node, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + if node == "w1" { + return workerctl.RunReply{}, cause + } + return workerctl.RunReply{}, nil + } + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).To(Succeed()) + Expect(control.seen()).To(HaveLen(2)) + }) + } + + It("returns an error that says nothing about the work when no worker took it", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + return workerctl.RunReply{}, workerctl.ErrNoRoute + } + err := handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil) + Expect(errors.Is(err, errNoAnswerYet)).To(BeTrue()) + Expect(control.seen()).To(HaveLen(3)) + }) + }) + + Describe("a carrier that is released", func() { + released := func() context.Context { + c, cancel := context.WithCancelCause(ctx) + cancel(messaging.ErrCarrierReleased) + return c + } + + It("releases a claim whose run never started, when no worker could be picked", func() { + picker.err = errors.New("no agent worker holds a tunnel") + err := handler(messaging.WorkMCPCI)(released(), ciPayload("j1"), nil) + Expect(err).To(HaveOccurred(), "completing it would drop work that was never run") + Expect(errors.Is(err, errNoAnswerYet)).To(BeTrue()) + Expect(control.seen()).To(BeEmpty()) + }) + + It("releases a claim whose worker was busy as the carrier was released", func() { + c, cancel := context.WithCancelCause(ctx) + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + cancel(messaging.ErrCarrierReleased) + return workerctl.RunReply{}, workerctl.ErrWorkerBusy + } + err := handler(messaging.WorkMCPCI)(c, ciPayload("j1"), nil) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, errNoAnswerYet)).To(BeTrue()) + }) + + It("completes the claim of a run that had started", func() { + c, cancel := context.WithCancelCause(ctx) + control.do = func(_ context.Context, _, _ string, onProgress func(string, json.RawMessage)) (workerctl.RunReply, error) { + onProgress("", json.RawMessage(`{}`)) + cancel(messaging.ErrCarrierReleased) + return workerctl.RunReply{}, context.Cause(c) + } + Expect(handler(messaging.WorkMCPCI)(c, ciPayload("j1"), nil)).To(Succeed()) + Expect(store.terminals()).To(BeEmpty(), "the reaper decides about the job") + }) + }) + + Describe("a job that was cancelled before it ran", func() { + It("does not dispatch a job whose row says cancelled, and completes the claim", func() { + store.statuses = map[string]string{"j1": "cancelled"} + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed()) + Expect(control.seen()).To(BeEmpty()) + }) + + It("does not start the job again when the cancel arrives while the worker is busy", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + Expect(bus.Publish(messaging.SubjectJobCancel("j1"), CancelEvent{JobID: "j1"})).To(Succeed()) + return workerctl.RunReply{}, workerctl.ErrWorkerBusy + } + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed(), "released it would run again") + Expect(control.seen()).To(HaveLen(1)) + Expect(store.terminals()).To(Equal([]terminal{{"j1", "cancelled", "", "cancelled"}})) + }) + + It("treats a pick that fails with the context ended by a cancel as a cancel, and not as a missing worker", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + Expect(bus.Publish(messaging.SubjectJobCancel("j2"), CancelEvent{JobID: "j2"})).To(Succeed()) + picker.mu.Lock() + picker.err = context.Canceled + picker.mu.Unlock() + return workerctl.RunReply{}, workerctl.ErrWorkerBusy + } + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j2"), nil)).To(Succeed()) + Expect(store.terminals()).To(Equal([]terminal{{"j2", "cancelled", "", "cancelled"}})) + }) + }) + Describe("the cancel of a job", func() { running := func() (started chan struct{}, finished chan error) { started = make(chan struct{}) diff --git a/core/services/jobs/dispatch_loop.go b/core/services/jobs/dispatch_loop.go index dc9790051568..e888d54cd369 100644 --- a/core/services/jobs/dispatch_loop.go +++ b/core/services/jobs/dispatch_loop.go @@ -40,6 +40,9 @@ type DispatchConfig struct { // MaxConcurrent is how many claims of a kind this replica drives at once, // DefaultMaxConcurrent when zero. MaxConcurrent int + // MaxFailures is how many failures end a queued unit that no worker can run. + // DefaultClaimMaxFailures when zero. + MaxFailures int } // DispatchLoop is what a frontend replica runs on the tunnel carrier: it claims @@ -70,6 +73,7 @@ func NewDispatchLoop(cfg DispatchConfig) (*DispatchLoop, error) { consumer, err := NewClaimConsumer(ClaimConsumerConfig{ DB: cfg.DB, Owner: cfg.Owner, Events: cfg.Hints, Hints: cfg.Hints, Interval: cfg.Interval, Liveness: cfg.Liveness, + MaxFailures: cfg.MaxFailures, OnExhausted: driver.failExhausted, }) if err != nil { return nil, err diff --git a/core/services/jobs/dispatch_loop_test.go b/core/services/jobs/dispatch_loop_test.go index 630b6ff6ad35..9fde6acc1190 100644 --- a/core/services/jobs/dispatch_loop_test.go +++ b/core/services/jobs/dispatch_loop_test.go @@ -123,10 +123,13 @@ var _ = Describe("The dispatch loop", func() { Describe("when the carrier of the loop is released", func() { // run starts a loop whose context the spec can end with a cause, and waits // until a run is in flight on it. - run := func() (cancel context.CancelCauseFunc, loop *DispatchLoop) { + run := func(sendLine bool) (cancel context.CancelCauseFunc, loop *DispatchLoop) { GinkgoHelper() started := make(chan struct{}, 1) - control.do = func(ctx context.Context, _, _ string, _ func(string, json.RawMessage)) (workerctl.RunReply, error) { + control.do = func(ctx context.Context, _, _ string, onProgress func(string, json.RawMessage)) (workerctl.RunReply, error) { + if sendLine { + onProgress("", json.RawMessage(`{"tick":1}`)) + } started <- struct{}{} <-ctx.Done() return workerctl.RunReply{}, ctx.Err() @@ -145,8 +148,8 @@ var _ = Describe("The dispatch loop", func() { return cancel, loop } - It("completes the claim of a run that the end of the drain cut off, and leaves its job to the reaper, so the run is not offered again", func() { - cancel, loop := run() + It("completes the claim of a run that the end of the drain cut off after it started, and leaves its job to the reaper, so the run is not offered again", func() { + cancel, loop := run(true) cancel(messaging.ErrCarrierReleased) loop.Stop() @@ -155,8 +158,18 @@ var _ = Describe("The dispatch loop", func() { Expect(control.seen()).To(HaveLen(1)) }) + It("releases the claim of a run that the end of the drain cut off before a line came back, so the carrier that takes over runs it", func() { + cancel, loop := run(false) + cancel(messaging.ErrCarrierReleased) + loop.Stop() + + var row WorkClaim + Expect(db.First(&row).Error).To(Succeed(), "the work was never run, so it must not be dropped") + Expect(row.State).To(Equal(ClaimPending)) + }) + It("still releases the claim of a run that ends for another reason, such as the stop of the replica", func() { - cancel, loop := run() + cancel, loop := run(false) cancel(context.Canceled) loop.Stop() @@ -184,6 +197,29 @@ var _ = Describe("The dispatch loop", func() { Expect(store.terminals()).To(Equal([]terminal{{"j1", "cancelled", "", "cancelled"}})) }) + It("fails the job and deletes the row of a unit that no worker can run, after the failures it allows", func() { + control.do = func(context.Context, string, string, func(string, json.RawMessage)) (workerctl.RunReply, error) { + // The wait of a release is not what this spec is about. + _ = db.Exec(`UPDATE work_claims SET not_before = NULL`).Error + return workerctl.RunReply{}, errors.New("the worker rejected the payload") + } + live(db, "replica-a") + loop, err := NewDispatchLoop(DispatchConfig{ + DB: db, Owner: "replica-a", Picker: picker, Control: control, Broadcast: broadcast, Store: store, Hints: bus, + Interval: 50 * time.Millisecond, MaxFailures: 2, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(loop.Start(ctx)).To(Succeed()) + DeferCleanup(loop.Stop) + Expect(NewClaimQueue(db, bus).Enqueue(ctx, messaging.WorkMCPCI, ci("j9"))).To(Succeed()) + + Eventually(rows, 15*time.Second).Should(BeZero()) + Expect(store.terminals()).To(HaveLen(1)) + Expect(store.terminals()[0].jobID).To(Equal("j9")) + Expect(store.terminals()[0].status).To(Equal("failed")) + Expect(store.terminals()[0].errMsg).To(ContainSubstring("repeated failures")) + }) + It("is built only with what it needs", func() { good := DispatchConfig{DB: db, Owner: "replica-a", Picker: picker, Control: control, Hints: bus} for name, mutate := range map[string]func(*DispatchConfig){ diff --git a/core/services/jobs/dispatcher.go b/core/services/jobs/dispatcher.go index 6f0a8f568b18..e6ca4d0b446a 100644 --- a/core/services/jobs/dispatcher.go +++ b/core/services/jobs/dispatcher.go @@ -194,8 +194,16 @@ func (d *Dispatcher) Enqueue(jobID, taskID, userID string) error { return d.queue.Enqueue(context.Background(), kind, evt) } -// Cancel publishes a cancel event to NATS (broadcast to all instances). +// Cancel marks the job as cancelled, and then publishes a cancel event to every +// replica. The mark is what stops a job that still waits in the queue: no replica +// holds a run to end, and the one that claims the unit reads the status before it +// dispatches. A job that already has an end state keeps it. func (d *Dispatcher) Cancel(jobID string) error { + if d.store != nil { + if err := d.store.UpdateJobStatus(jobID, "cancelled", "", "cancelled"); err != nil { + return fmt.Errorf("marking job %q as cancelled: %w", jobID, err) + } + } return d.nats.Publish(messaging.SubjectJobCancel(jobID), CancelEvent{ JobID: jobID, }) diff --git a/core/services/jobs/dispatcher_test.go b/core/services/jobs/dispatcher_test.go index ccf0c1720ec6..386bb1c249bf 100644 --- a/core/services/jobs/dispatcher_test.go +++ b/core/services/jobs/dispatcher_test.go @@ -489,3 +489,39 @@ var _ = Describe("Dispatcher", func() { type broadcastOnly struct { messaging.Broadcaster } + +var _ = Describe("Dispatcher.Cancel", func() { + var ( + store *JobStore + disp *Dispatcher + bus *testutil.FakeBus + ) + + BeforeEach(func() { + db := testutil.SetupTestDB() + var err error + store, err = NewJobStore(db) + Expect(err).ToNot(HaveOccurred()) + bus = testutil.NewFakeBus() + disp = NewDispatcher(store, &fakeWorkQueue{}, bus, db, "test-instance") + Expect(store.CreateTask(&TaskRecord{ID: "t1", UserID: "u1"})).To(Succeed()) + }) + + It("marks a job that waits in the queue as cancelled, so that no replica runs it", func() { + Expect(store.CreateJob(&JobRecord{ID: "j1", TaskID: "t1", UserID: "u1", Status: "pending"})).To(Succeed()) + + Expect(disp.Cancel("j1")).To(Succeed()) + + status, err := store.JobStatus("j1") + Expect(err).ToNot(HaveOccurred()) + Expect(status).To(Equal("cancelled")) + }) + + It("keeps the end state of a job that has finished", func() { + Expect(store.CreateJob(&JobRecord{ID: "j2", TaskID: "t1", UserID: "u1", Status: "completed"})).To(Succeed()) + Expect(disp.Cancel("j2")).To(Succeed()) + status, err := store.JobStatus("j2") + Expect(err).ToNot(HaveOccurred()) + Expect(status).To(Equal("completed")) + }) +}) diff --git a/core/services/jobs/store.go b/core/services/jobs/store.go index 2adf2da8d88c..668cfc51d9dd 100644 --- a/core/services/jobs/store.go +++ b/core/services/jobs/store.go @@ -213,6 +213,15 @@ func (s *JobStore) DeleteJob(id string) error { return s.db.Where("id = ?", id).Delete(&JobRecord{}).Error } +// JobStatus returns the status of a job. +func (s *JobStore) JobStatus(id string) (string, error) { + var status string + if err := s.db.Model(&JobRecord{}).Select("status").Where("id = ?", id).Scan(&status).Error; err != nil { + return "", err + } + return status, nil +} + // UpdateJobStatus updates just the status (and optionally result/error) of a job. func (s *JobStore) UpdateJobStatus(id, status, result, errMsg string) error { updates := map[string]any{ diff --git a/core/services/nodes/control_nats.go b/core/services/nodes/control_nats.go index eb79dfd99181..fe26b8869981 100644 --- a/core/services/nodes/control_nats.go +++ b/core/services/nodes/control_nats.go @@ -27,7 +27,7 @@ import ( // // It is NOT returned for a timeout, for a transport fault, or for a worker that // answered with a refusal. A node that answers is present by demonstration. -var ErrNoRoute = errors.New("nodes: no route to that node") +var ErrNoRoute = workerctl.ErrNoRoute // natsLink carries control requests as NATS requests to the per-node subject of // the verb. It maps the failure of the carrier onto the conditions this package diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index cbfe6a8873e7..d11b53796b82 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -659,7 +659,7 @@ func (a *nodeControl) DeleteModelFiles(modelName string) error { // It is a fact about the version of the worker. It is not a routing fact and // not an answer about any backend, so only a caller that has a fallback for an // older worker reads it, and it becomes ErrNoRoute there. -var errVerbNotServed = errors.New("the worker does not serve that control verb") +var errVerbNotServed = workerctl.ErrVerbNotServed // callVerb sends one request and returns the decoded reply, or nil and the // error. diff --git a/core/services/workerctl/request.go b/core/services/workerctl/request.go index 52466c4f3ddf..90e3dfa92dc5 100644 --- a/core/services/workerctl/request.go +++ b/core/services/workerctl/request.go @@ -66,6 +66,16 @@ const BusyMessage = "the worker has no free slot for this run" // route: the route exists. var ErrWorkerBusy = errors.New("the worker has no free slot") +// ErrNoRoute means that nothing answers for a node: there is no route to it now. +// It is a routing fact and says nothing about the node or about the work. Callers +// outside the nodes package, such as the claim queue, read it here. +var ErrNoRoute = errors.New("nodes: no route to that node") + +// ErrVerbNotServed means that a worker answered and does not serve the control +// verb. It is a fact about the version of the worker: a run that was sent to it +// did not start. +var ErrVerbNotServed = errors.New("the worker does not serve that control verb") + // WriteBusy answers a run that this worker cannot take now. The worker is // present and routable, and the request was not started, so the frontend may // offer it to another worker. From 70c58acee45db00c8065c815ff211e0a80c1632e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:10:57 +0000 Subject: [PATCH 137/162] fix(carrier): put queued work back when the new carrier refuses it, and purge what moved MigratePending marks the rows as migrated before they are published. A row that the new carrier refused stayed migrated and was lost. It goes back to pending now, so the next sweep takes it, and only the refused rows go back. Rows that moved stay for a day as evidence and are then deleted with their payloads, in the handoff and on the poll of a consumer. The seams document now says what the claim queue stores: the payload of a row is the whole request of the work, including the configuration of the MCP servers of a job, in plain text until the work completes or the row is purged. It also states the release rules, the failure cap and the cancel path. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .agents/distributed-seams.md | 54 +++++++-- core/services/carrier/claims.go | 21 +++- core/services/carrier/claims_handoff_test.go | 110 +++++++++++++++++++ core/services/carrier/export_test.go | 6 + 4 files changed, 182 insertions(+), 9 deletions(-) create mode 100644 core/services/carrier/claims_handoff_test.go create mode 100644 core/services/carrier/export_test.go diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index c380b3362070..5957062ac966 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -284,9 +284,20 @@ handler. wait keeps one poison row from being claimed ahead of every newer row on every tick). An error that wraps `jobs.ErrKeepClaim` leaves it with its replica: the work ran and its answer could not be recorded, and releasing would run it again. -- There is no dead letter. The only outcomes that release a row are those where - nothing was learned about the work, and failing a job on them would report a - missing connection as the verdict of a worker. +- A release is for work that was not started. The driver releases a claim when no + worker could be picked, when a worker has no slot, no route or is too old to + serve the verb (the next worker is tried first, with no wait), when the carrier + was released before the run began, and when a worker was reached and the call + failed before one line of the run came back. Once a line of the run has come + back, the claim is never released: a broken link then closes the job as failed + and completes the claim, because a second run could repeat a CI pipeline or the + side effects of an agent. +- A row ends after `MaxFailures` failures (default 10, about five minutes of + trying; `DispatchConfig.MaxFailures`). Only the failures that say something about + the work count: a worker was reached and the run failed on it. A release for lack + of a worker, a route or a free slot does not count, so a fleet that is down for a + day keeps its queue. The consumer records the job as failed first and deletes the + row after. - A claim is settled only by the replica and the attempt that hold it, so a holder that was reaped and answers late cannot delete the claim of the replica that took the work over. @@ -294,16 +305,43 @@ handler. `cluster.LiveInstanceIDsSQL` and has no age in it: a live replica that holds a claim for an hour has a slow job. A replica whose own row is not live claims nothing. +- A replica stops the handler of a claim that is no longer its own. Every tick the + consumer (once for all kinds) checks that the replica is live and lists the + claims it holds. When the instances row is missing for as long as a peer needs to + reap, the handler's context is cancelled, the peer runs the work, and this + replica does not settle the row. When this replica registers again before a peer + reaped, it gives the row back to the queue itself, because a row held by a live + id is never reaped. +- A process that starts to claim first releases the rows its own id holds + (`ReleaseOwned`). `LOCALAI_INSTANCE_ID` is stable across restarts, the id is + live again at once, and the reap never frees rows of a live id. - Delivery is at-least-once. A replica that dies in the middle of a job loses - nothing and the job runs again. NATS loses the job of a consumer that dies. + nothing and the job runs again. NATS loses the job of a consumer that dies. A + stall of a replica for as long as the liveness window can also run a job twice. + A handler for work that is not safe to repeat must record its progress. - `MigratePending` takes the pending rows out of the queue once, for a change to a carrier that has its own queue. Held rows are driven to their end where they are. + A row the new carrier refuses goes back to pending (`RestorePending`). Moved rows + stay as `migrated` for `MigratedRetention` (a day) and are then deleted with + their payloads. +- Payload at rest: the payload of a row is the whole request of the work, in + plain text. For an MCP job it holds the job, the task and the model + configuration, including the configuration of its MCP servers. It stays in + `work_claims.payload` until the work is complete (the row is deleted) or, for a + moved row, until the purge. Anyone who can read the table can read it. Treat the + table as the place where those secrets live, and give the database role the + same care. +- The poll reads the rows of a kind through `idx_work_claims_pick` (kind, state, + created_at) and filters the rows that wait for a backoff, so it passes over them + at each poll. The failure cap keeps their number small. - A plain task (`WorkTask`) has no worker on either carrier. The driver drops the claim and the job is left to the reaper (`ReapStuckJobs`), as it was on NATS. -- `Dispatcher.Cancel` publishes `jobs..cancel`. On NATS nothing subscribes, as - on master, and a spec pins that. On this carrier the replica that holds the job - ends its stream, the run on the worker ends with its request, and the job is - closed as cancelled. +- `Dispatcher.Cancel` marks the job as cancelled in the database and publishes + `jobs..cancel`. The mark is what stops a job that is still queued: the driver + reads the status before it dispatches, and again before each worker it offers the + run to. On NATS nothing subscribes to the subject, as on master, and a spec pins + that. On this carrier the replica that holds the job ends its stream, the run on + the worker ends with its request, and the job is closed as cancelled. ## Control verbs diff --git a/core/services/carrier/claims.go b/core/services/carrier/claims.go index 935cf7255069..65406fedffd7 100644 --- a/core/services/carrier/claims.go +++ b/core/services/carrier/claims.go @@ -64,6 +64,12 @@ func (c *ClaimWork) start(set func() *Set, selector *nodes.AgentSelector, token // queue of the carrier that took over. The rows are taken out of the table first // and then published, so a job is lost, and not run twice, if this replica dies in // between. A job that is lost stays running until the reaper fails it. +// +// A row that the new carrier refuses goes back to pending, so that nothing is +// lost while the new carrier is down: the rows wait in the table, and the next +// sweep (the handoff runs again after a pause) or the next period on the tunnel +// takes them. The rows that were moved stay as "migrated" for a day and are then +// purged with their payloads. func (c *ClaimWork) handoff() func(context.Context, *Set) error { return func(ctx context.Context, next *Set) error { rows, err := jobs.MigratePending(ctx, c.DB) @@ -71,16 +77,29 @@ func (c *ClaimWork) handoff() func(context.Context, *Set) error { return err } var failed error + var refused []string moved := 0 for _, r := range rows { if err := next.WorkQueue.Enqueue(ctx, messaging.WorkKind(r.Kind), json.RawMessage(r.Payload)); err != nil { failed = errors.Join(failed, fmt.Errorf("publishing the %s unit %s: %w", r.Kind, r.ID, err)) + refused = append(refused, r.ID) continue } moved++ } + if len(refused) > 0 { + // Not bound to ctx: the reason for the failure may be that it ended. + restoreCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() + if err := jobs.RestorePending(restoreCtx, c.DB, refused); err != nil { + failed = errors.Join(failed, fmt.Errorf("putting %d refused units back in the queue: %w", len(refused), err)) + } + } if len(rows) > 0 { - xlog.Info("Queued work moved to the carrier that took over", "carrier", next.Name, "moved", moved, "lost", len(rows)-moved) + xlog.Info("Queued work moved to the carrier that took over", "carrier", next.Name, "moved", moved, "kept_in_queue", len(refused)) + } + if _, err := jobs.PurgeMigrated(ctx, c.DB, jobs.MigratedRetention); err != nil { + failed = errors.Join(failed, err) } return failed } diff --git a/core/services/carrier/claims_handoff_test.go b/core/services/carrier/claims_handoff_test.go new file mode 100644 index 000000000000..f1b18401a66f --- /dev/null +++ b/core/services/carrier/claims_handoff_test.go @@ -0,0 +1,110 @@ +package carrier_test + +import ( + "context" + "errors" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/jobs" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// recordingQueue is the queue of the carrier that takes over. +type recordingQueue struct { + mu sync.Mutex + err error + seen []messaging.WorkKind +} + +func (q *recordingQueue) Enqueue(_ context.Context, kind messaging.WorkKind, _ any) error { + q.mu.Lock() + defer q.mu.Unlock() + if q.err != nil { + return q.err + } + q.seen = append(q.seen, kind) + return nil +} + +var _ = Describe("The handoff of queued work to the carrier that takes over", func() { + var ( + ctx context.Context + work *carrier.ClaimWork + next *carrier.Set + q *recordingQueue + ids []string + ) + + count := func(state string) int64 { + GinkgoHelper() + var n int64 + Expect(work.DB.Model(&jobs.WorkClaim{}).Where("state = ?", state).Count(&n).Error).To(Succeed()) + return n + } + + BeforeEach(func() { + ctx = context.Background() + db := testutil.SetupTestDB() + Expect(jobs.MigrateClaims(ctx, db)).To(Succeed()) + work = &carrier.ClaimWork{DB: db, Owner: "replica-a"} + q = &recordingQueue{} + next = &carrier.Set{Name: "nats", WorkQueue: q} + ids = nil + for _, kind := range []messaging.WorkKind{messaging.WorkMCPCI, messaging.WorkAgentRun} { + id, err := jobs.EnqueueClaim(ctx, db, kind, map[string]string{"k": string(kind)}) + Expect(err).ToNot(HaveOccurred()) + ids = append(ids, id) + } + }) + + It("puts the rows back in the queue when the new carrier refuses them, and loses none", func() { + q.err = errors.New("the broker is down") + + err := work.HandoffForTest()(ctx, next) + + Expect(err).To(MatchError(ContainSubstring("the broker is down"))) + Expect(count(jobs.ClaimPending)).To(BeEquivalentTo(2), "both rows wait in the table") + Expect(count(jobs.ClaimMigrated)).To(BeZero()) + + // The broker is back: the next sweep moves them. + q.err = nil + Expect(work.HandoffForTest()(ctx, next)).To(Succeed()) + Expect(q.seen).To(ConsistOf(messaging.WorkMCPCI, messaging.WorkAgentRun)) + Expect(count(jobs.ClaimMigrated)).To(BeEquivalentTo(2)) + Expect(count(jobs.ClaimPending)).To(BeZero()) + }) + + It("restores only the rows that were refused", func() { + q2 := &selectiveQueue{refuse: messaging.WorkAgentRun} + next.WorkQueue = q2 + Expect(work.HandoffForTest()(ctx, next)).ToNot(Succeed()) + Expect(count(jobs.ClaimMigrated)).To(BeEquivalentTo(1)) + Expect(count(jobs.ClaimPending)).To(BeEquivalentTo(1)) + }) + + It("purges the rows it moved, with their payloads, after the retention", func() { + Expect(work.HandoffForTest()(ctx, next)).To(Succeed()) + Expect(count(jobs.ClaimMigrated)).To(BeEquivalentTo(2)) + + Expect(work.DB.Exec(`UPDATE work_claims SET claimed_at = now() - interval '25 hours'`).Error).To(Succeed()) + Expect(work.HandoffForTest()(ctx, next)).To(Succeed()) + Expect(count(jobs.ClaimMigrated)).To(BeZero()) + var left int64 + Expect(work.DB.Model(&jobs.WorkClaim{}).Where("id IN ?", ids).Count(&left).Error).To(Succeed()) + Expect(left).To(BeZero(), "no row, and so no payload, is left behind") + }) +}) + +type selectiveQueue struct{ refuse messaging.WorkKind } + +func (q *selectiveQueue) Enqueue(_ context.Context, kind messaging.WorkKind, _ any) error { + if kind == q.refuse { + return errors.New("refused") + } + return nil +} diff --git a/core/services/carrier/export_test.go b/core/services/carrier/export_test.go new file mode 100644 index 000000000000..2b01750ccd9a --- /dev/null +++ b/core/services/carrier/export_test.go @@ -0,0 +1,6 @@ +package carrier + +import "context" + +// HandoffForTest exposes the Handoff of the tunnel set to the specs. +func (c *ClaimWork) HandoffForTest() func(context.Context, *Set) error { return c.handoff() } From 6e4c9cb1a40d47f042806a7c22d29658c5c37cfd Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:12:51 +0000 Subject: [PATCH 138/162] fix(nodes): bind a worker's broadcasts to its run, and require a control token A worker may ask the frontend to publish a progress line on a subject. The allow list of the node type was the only limit, so a worker that held one job could publish progress, results or agent events of any other job or user, and a subject with an empty token or white space reached the filters. - The driver now allows a run only the subjects of the run: the progress and the result of its job, or the events of its agent and user. A run that names no scope keeps the allow list of its node type. - A subject with an empty token, white space or a control character is refused before the filters see it. - agentworker.Start refuses an empty control token. The control server listens on loopback, and an empty token left it open to every process of the host. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/agentworker/start.go | 11 +++++-- core/services/agentworker/start_test.go | 13 ++++++-- core/services/jobs/claim_driver.go | 42 +++++++++++++++++++++++-- core/services/jobs/claim_driver_test.go | 38 ++++++++++++++++++++++ core/services/nodes/rebroadcast.go | 21 ++++++++++++- core/services/nodes/rebroadcast_test.go | 8 +++++ 6 files changed, 124 insertions(+), 9 deletions(-) diff --git a/core/services/agentworker/start.go b/core/services/agentworker/start.go index e8f6df8575da..91e8fd789957 100644 --- a/core/services/agentworker/start.go +++ b/core/services/agentworker/start.go @@ -29,9 +29,9 @@ type Options struct { // credential. It may be nil. Reauthorize func(ctx context.Context) error // ControlToken guards the control server. It is the registration token of - // the deployment, which the control client of the frontend presents. An empty - // token leaves the server open to anything that can reach it, which is the - // tunnel and this host. + // the deployment, which the control client of the frontend presents. It is + // required: an empty token would leave the loopback port open to every process + // of this host. ControlToken string // Handler serves the verbs. See Handler. Handler http.Handler @@ -92,6 +92,11 @@ func Start(ctx context.Context, o Options) (*Runtime, error) { if o.TunnelToken == nil { return nil, errors.New("starting the agent worker control plane: no way to read the tunnel credential") } + if o.ControlToken == "" { + return nil, errors.New("starting the agent worker control plane: no control token. " + + "Set LOCALAI_REGISTRATION_TOKEN to the registration token of the deployment, on the worker and on the frontends; " + + "without it any process on this host could call the control port") + } lis, err := net.Listen("tcp", loopbackBind) if err != nil { return nil, fmt.Errorf("binding the agent worker control server: %w", err) diff --git a/core/services/agentworker/start_test.go b/core/services/agentworker/start_test.go index a0abf9ee0a9c..fb179e310bed 100644 --- a/core/services/agentworker/start_test.go +++ b/core/services/agentworker/start_test.go @@ -80,15 +80,24 @@ var _ = Describe("Starting the control plane of an agent worker", func() { }) It("refuses to start with no handler or no token function", func() { - _, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{FrontendURL: "http://127.0.0.1:1", NodeID: "n", TunnelToken: func() string { return "t" }}) + _, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{FrontendURL: "http://127.0.0.1:1", NodeID: "n", TunnelToken: func() string { return "t" }, ControlToken: "secret"}) Expect(err).To(HaveOccurred()) - _, err = agentworker.Start(GinkgoT().Context(), agentworker.Options{FrontendURL: "http://127.0.0.1:1", NodeID: "n", Handler: http.NotFoundHandler()}) + _, err = agentworker.Start(GinkgoT().Context(), agentworker.Options{FrontendURL: "http://127.0.0.1:1", NodeID: "n", Handler: http.NotFoundHandler(), ControlToken: "secret"}) Expect(err).To(HaveOccurred()) }) + It("refuses to start with no control token, because the loopback port would be open to any process of the host", func() { + _, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{ + FrontendURL: "http://127.0.0.1:1", NodeID: "n", TunnelToken: func() string { return "t" }, + Handler: http.NotFoundHandler(), + }) + Expect(err).To(MatchError(ContainSubstring("LOCALAI_REGISTRATION_TOKEN"))) + }) + It("refuses a frontend URL that cannot be dialled", func() { _, err := agentworker.Start(GinkgoT().Context(), agentworker.Options{ FrontendURL: "ftp://frontend", NodeID: "n", TunnelToken: func() string { return "t" }, Handler: http.NotFoundHandler(), + ControlToken: "secret", }) Expect(err).To(HaveOccurred()) }) diff --git a/core/services/jobs/claim_driver.go b/core/services/jobs/claim_driver.go index 7044d68405ce..26ec07c4513a 100644 --- a/core/services/jobs/claim_driver.go +++ b/core/services/jobs/claim_driver.go @@ -203,6 +203,8 @@ func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb s defer d.cancelled.remove(jobID) } + scope := broadcastScope(kind, jobID, payload) + // started becomes true when the first line of the run comes back. From then on // the job may have done work, and the claim is not released. var started atomic.Bool @@ -232,7 +234,7 @@ func (d *AgentDriver) drive(ctx context.Context, kind messaging.WorkKind, verb s err = d.cfg.Control.CallStreaming(ctx, nodeID, verb, json.RawMessage(payload), &reply, func(subject string, raw json.RawMessage) { started.Store(true) - d.rebroadcast(nodeType, subject, raw) + d.rebroadcast(nodeType, subject, raw, scope) }) switch { case err == nil: @@ -332,14 +334,48 @@ func (d *AgentDriver) cutOff(jobID string) error { return nil } +// broadcastScope returns the subjects that the run being driven may publish on, or +// nil when the run names none. A worker that holds one run may ask the frontend to +// publish the progress of that run, and not the progress or the events of another +// job or another user: the allow list of the node type is wide, and the run is the +// narrower right. +func broadcastScope(kind messaging.WorkKind, jobID string, payload []byte) map[string]bool { + switch kind { + case messaging.WorkMCPCI: + if jobID == "" { + return nil + } + return map[string]bool{ + messaging.SubjectJobProgress(jobID): true, + messaging.SubjectJobResult(jobID): true, + } + case messaging.WorkAgentRun: + var evt struct { + AgentName string `json:"agent_name"` + UserID string `json:"user_id"` + } + if err := json.Unmarshal(payload, &evt); err != nil || evt.AgentName == "" { + return nil + } + return map[string]bool{messaging.SubjectAgentEvents(evt.AgentName, evt.UserID): true} + } + return nil +} + // rebroadcast publishes the broadcast that a progress line asked for. A line with // no subject is meant for the caller alone, as every private tick is, and is not // passed on: the allow list denies the empty subject, and passing it would log a -// refusal for each ordinary tick. -func (d *AgentDriver) rebroadcast(nodeType, subject string, raw json.RawMessage) { +// refusal for each ordinary tick. A line that names a subject outside the scope of +// the run is refused. A run with no scope is held to the allow list of its node +// type only. +func (d *AgentDriver) rebroadcast(nodeType, subject string, raw json.RawMessage, scope map[string]bool) { if subject == "" || d.cfg.Broadcast == nil { return } + if scope != nil && !scope[subject] { + xlog.Warn("Refusing a broadcast that a worker asked for outside the run it was given", "nodeType", nodeType, "subject", subject) + return + } d.cfg.Broadcast.Handle(nodeType, subject, raw) } diff --git a/core/services/jobs/claim_driver_test.go b/core/services/jobs/claim_driver_test.go index f4a83327e1d3..b2f1bc8cc0c3 100644 --- a/core/services/jobs/claim_driver_test.go +++ b/core/services/jobs/claim_driver_test.go @@ -282,6 +282,44 @@ var _ = Describe("The agent driver", func() { }) }) + Describe("the broadcasts of a run", func() { + lines := func(subjects ...string) { + control.do = func(_ context.Context, _, _ string, onProgress func(string, json.RawMessage)) (workerctl.RunReply, error) { + for _, s := range subjects { + onProgress(s, json.RawMessage(`{}`)) + } + return workerctl.RunReply{}, nil + } + } + published := func() []string { + broadcast.mu.Lock() + defer broadcast.mu.Unlock() + var out []string + for _, l := range broadcast.lines { + out = append(out, l.subject) + } + return out + } + + It("publishes the progress and the result of the job that it drives, and no other job", func() { + lines("jobs.j1.progress", "jobs.j1.result", "jobs.other.progress", "jobs.other.result", "agent.a.events.u") + Expect(handler(messaging.WorkMCPCI)(ctx, ciPayload("j1"), nil)).To(Succeed()) + Expect(published()).To(Equal([]string{"jobs.j1.progress", "jobs.j1.result"})) + }) + + It("publishes the events of the agent and the user that it drives, and no other", func() { + lines("agent.helper.events.alice", "agent.helper.events.bob", "agent.other.events.alice", "jobs.j1.progress") + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{"agent_name":"helper","user_id":"alice"}`), nil)).To(Succeed()) + Expect(published()).To(Equal([]string{"agent.helper.events.alice"})) + }) + + It("holds a run that names no scope to the allow list of its node type", func() { + lines("jobs.x.progress") + Expect(handler(messaging.WorkAgentRun)(ctx, []byte(`{}`), nil)).To(Succeed()) + Expect(published()).To(Equal([]string{"jobs.x.progress"})) + }) + }) + Describe("a run whose link breaks", func() { // startedThen plays a worker that sends a line of the run and then loses the link. startedThen := func(err error) { diff --git a/core/services/nodes/rebroadcast.go b/core/services/nodes/rebroadcast.go index 9a3172d8ed7e..69d490863f3b 100644 --- a/core/services/nodes/rebroadcast.go +++ b/core/services/nodes/rebroadcast.go @@ -3,6 +3,7 @@ package nodes import ( "encoding/json" "strings" + "unicode" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/xlog" @@ -37,6 +38,24 @@ func mayBroadcast(nodeType, subject string) bool { return mayBroadcastIn(workerBroadcastAllow, nodeType, subject) } +// wellFormedSubject says whether every token of subject is not empty and holds +// no white space and no control character. The filters of the table match a +// token by position, so a subject with an empty token, or one that a log line or +// a subscriber would read as two, must not reach them. +func wellFormedSubject(subject string) bool { + for _, token := range strings.Split(subject, ".") { + if token == "" { + return false + } + for _, r := range token { + if unicode.IsSpace(r) || unicode.IsControl(r) { + return false + } + } + } + return true +} + // mayBroadcastIn is mayBroadcast against a caller-supplied table. // // It exists so a spec can hold the table itself constant while varying what is @@ -48,7 +67,7 @@ func mayBroadcastIn(table map[string][]string, nodeType, subject string) bool { // is nothing here to authorise. Refusing it is not defensive tidiness: it // keeps a filter that happened to match the empty string from turning every // ordinary private progress tick into a publish. - if subject == "" || strings.ContainsAny(subject, "*>") { + if subject == "" || strings.ContainsAny(subject, "*>") || !wellFormedSubject(subject) { // A wildcard is a filter for a subscriber. A worker that names one asks for // a publish on every subject it matches, and a filter of the table would // match its own text. diff --git a/core/services/nodes/rebroadcast_test.go b/core/services/nodes/rebroadcast_test.go index b0a7734ce935..c3ff192dfccd 100644 --- a/core/services/nodes/rebroadcast_test.go +++ b/core/services/nodes/rebroadcast_test.go @@ -52,6 +52,14 @@ var _ = Describe("The broadcasts that a worker may ask for", func() { Entry("may not publish on a root that carries state", "state.models.delta", false), Entry("may not publish on a wildcard", "jobs.*.result", false), Entry("may not publish on no subject", "", false), + Entry("may not publish with an empty token", "jobs..progress", false), + Entry("may not publish with an empty token at the end", "agent.helper.events.", false), + Entry("may not publish with an empty agent name", "agent..events.alice", false), + Entry("may not publish with a trailing dot", "jobs.j1.progress.", false), + Entry("may not publish with a space in a token", "jobs.j 1.progress", false), + Entry("may not publish with a newline in a token", "jobs.j1\n.result", false), + Entry("may not publish with a tab in a token", "agent.hel\tper.events.alice", false), + Entry("may not publish with a control character", "jobs.j\x001.result", false), ) It("denies a backend worker and an unknown type everything", func() { From 571a7a290001702d8ff24adc82e5f40b9b53abdd Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:19:11 +0000 Subject: [PATCH 139/162] fix(carrier): bound the stop and the hand-over of a swap, close the window gap, re-attach a drain Review of the swap found six robustness defects. - retire waited for the work of the old set without a bound, so one stuck handler held the hand-over of the queue and the close of the set. The wait is now bounded (StopGrace, 30 s), logged, and left behind. Close stops the sets in parallel under the same bound. - Close could block for two minutes in the hand-over, which ran on a context that nothing cancelled. The hand-over now runs on a context that Close cancels after CloseGrace (10 s). - swapTo stored the new set before it opened the window, and a call to a worker attached only to the old carrier failed in the gap. The window opens first. - The holder kept a reconnect hook entry for every set it had ever hooked. It drops the entry when the set is released, on abort and on close too. - A replica that started or restarted during a drain never attached the draining carrier, so workers attached only to it were not reached until the drain ended. It now builds that carrier again, opens the window and starts its work, and tries again at the next poll if it cannot. - A failure to start the work of the active carrier was logged and not reported. The replica now reports not ready with the reason, and ready again once the work starts. The fakes of the swap specs answer with an error once their set is closed, so that the spec "never reaches a closed set" can fail. It now also calls a worker through the window during the drain. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/broadcaster.go | 3 + core/services/carrier/export_test.go | 7 + core/services/carrier/fakes_test.go | 14 +- core/services/carrier/swap.go | 161 ++++++++++++++++++---- core/services/carrier/swap_test.go | 199 ++++++++++++++++++++++++++- 5 files changed, 354 insertions(+), 30 deletions(-) diff --git a/core/services/carrier/broadcaster.go b/core/services/carrier/broadcaster.go index ebf8702a6731..d1fc0c90151c 100644 --- a/core/services/carrier/broadcaster.go +++ b/core/services/carrier/broadcaster.go @@ -126,6 +126,9 @@ func (b *Broadcaster) Release(old *Set) error { return nil } b.listening = slices.Delete(b.listening, i, i+1) + // A released set is finished and is closed by the caller. The map must not keep + // it alive for the life of the holder. + delete(b.hooked, old) subs := b.snapshotLocked() b.mu.Unlock() diff --git a/core/services/carrier/export_test.go b/core/services/carrier/export_test.go index 2b01750ccd9a..9de967a7702e 100644 --- a/core/services/carrier/export_test.go +++ b/core/services/carrier/export_test.go @@ -4,3 +4,10 @@ import "context" // HandoffForTest exposes the Handoff of the tunnel set to the specs. func (c *ClaimWork) HandoffForTest() func(context.Context, *Set) error { return c.handoff() } + +// HookedCount is how many sets the holder keeps a reconnect hook for. +func (b *Broadcaster) HookedCount() int { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.hooked) +} diff --git a/core/services/carrier/fakes_test.go b/core/services/carrier/fakes_test.go index 06799a4dcdd7..d6307f779b24 100644 --- a/core/services/carrier/fakes_test.go +++ b/core/services/carrier/fakes_test.go @@ -35,7 +35,11 @@ func (r *recorder) seen() []string { } // fakeCommands is a nodes.NodeControl that only records the call. -type fakeCommands struct{ recorder } +type fakeCommands struct { + recorder + // closed, when set, makes the calls that a spec hammers fail once the set is closed. + closed interface{ Load() bool } +} var _ nodes.NodeControl = (*fakeCommands)(nil) @@ -63,7 +67,13 @@ func (f *fakeCommands) UnloadModelOnNode(string, string) error { f.rec("UnloadModelOnNode") return nil } -func (f *fakeCommands) PingNode(string) error { f.rec("PingNode"); return nil } +func (f *fakeCommands) PingNode(string) error { + if f.closed != nil && f.closed.Load() { + return errUsedAfterClose + } + f.rec("PingNode") + return nil +} func (f *fakeCommands) InstallBackendOp(string, string, string, string, int, string, string, time.Duration, func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { f.rec("InstallBackendOp") return &workerctl.BackendInstallReply{}, nil diff --git a/core/services/carrier/swap.go b/core/services/carrier/swap.go index 45ee91f29f05..e4cc74f0fd29 100644 --- a/core/services/carrier/swap.go +++ b/core/services/carrier/swap.go @@ -35,6 +35,12 @@ type ReadyReporter interface { ReportReady(ctx context.Context, epoch int64, reason string) error } +// WindowControl is the part of Window that the swapper moves. *Window is one. +type WindowControl interface { + Open(prev *Set) + Close() +} + // DefaultSwapInterval is how often a replica reads the cluster carrier row. const DefaultSwapInterval = 2 * time.Second @@ -45,7 +51,7 @@ type SwapperOptions struct { // Bus is the broadcaster holder over Cur. Bus *Broadcaster // Window routes the calls to workers during a change. It may be nil. - Window *Window + Window WindowControl Rows RowSource Ready ReadyReporter Build Builder @@ -55,7 +61,14 @@ type SwapperOptions struct { // Settle is how long after the release of a carrier its handoff runs again, // for the producers that were slow to flip. Twice the interval when zero. Settle time.Duration - Meter metric.Meter + // StopGrace bounds the wait for the work of a set to stop. A set whose work + // does not stop in time is logged and left behind, so that the hand-over and + // the close of the set are not held up by it. 30 seconds when zero. + StopGrace time.Duration + // CloseGrace bounds the wait of Close for the hand-over that runs in the + // background. After it the hand-over is cancelled. 10 seconds when zero. + CloseGrace time.Duration + Meter metric.Meter } // Swapper is what a replica does in a change of carrier. It follows the row: @@ -81,6 +94,7 @@ type Swapper struct { prepared *Set // built for the change in prepare, attached, not yet in use draining *Set // the previous set, attached through the drain members map[*Set]*member + startErr map[*Set]error // why the work of a set did not start, until it does reported struct { set bool epoch int64 @@ -89,6 +103,11 @@ type Swapper struct { prepareSeen time.Time closed bool + // life is the context of the hand-over that runs in the background. It ends + // when Close gives up waiting for it. + life context.Context + lifeCancel context.CancelFunc + wake chan struct{} done chan struct{} doneOnce sync.Once @@ -125,11 +144,19 @@ func NewSwapper(o SwapperOptions) (*Swapper, error) { if o.Settle <= 0 { o.Settle = 2 * o.Interval } + if o.StopGrace <= 0 { + o.StopGrace = 30 * time.Second + } + if o.CloseGrace <= 0 { + o.CloseGrace = 10 * time.Second + } + life, lifeCancel := context.WithCancel(context.Background()) s := &Swapper{ - o: o, base: context.Background(), - members: map[*Set]*member{}, - wake: make(chan struct{}, 1), - done: make(chan struct{}), + o: o, base: context.Background(), life: life, lifeCancel: lifeCancel, + members: map[*Set]*member{}, + startErr: map[*Set]error{}, + wake: make(chan struct{}, 1), + done: make(chan struct{}), } s.registerMetrics() return s, nil @@ -270,7 +297,12 @@ func (s *Swapper) Reconcile(ctx context.Context, row cluster.CarrierRow) error { return s.report(ctx, row.Epoch, fmt.Sprintf("cannot build the %s carrier: %v", row.Active, err)) } } - s.ensureStarted(s.o.Cur.Load()) + // A replica that cannot run the work of its carrier is not ready. It says so, + // and the next poll tries again. + if cur := s.o.Cur.Load(); s.ensureStarted(cur) != nil { + s.countFailure() + return s.report(ctx, row.Epoch, fmt.Sprintf("cannot start the work of the %s carrier: %v", cur.Name, s.startErr[cur])) + } switch row.State { case cluster.StatePrepare: @@ -279,7 +311,7 @@ func (s *Swapper) Reconcile(ctx context.Context, row cluster.CarrierRow) error { return s.report(ctx, row.Epoch, "") default: s.prepareSeen = time.Time{} - s.settle(row) + s.settle(ctx, row) return nil } } @@ -368,6 +400,14 @@ func (s *Swapper) swapTo(ctx context.Context, row cluster.CarrierRow) error { return err } old := s.o.Cur.Load() + // The window opens before the pointer moves. While it is open with the old set + // as both the current and the previous set, every call goes to the old set, + // which is where it went. Once the pointer moves, a call to a worker that is + // attached to the old carrier only finds the window already open. The other + // order leaves a gap in which such a call goes to the new carrier and fails. + if s.o.Window != nil { + s.o.Window.Open(old) + } s.o.Cur.Store(next) third := s.draining @@ -378,15 +418,12 @@ func (s *Swapper) swapTo(ctx context.Context, row cluster.CarrierRow) error { s.prepared = nil } s.draining = old - if s.o.Window != nil { - s.o.Window.Open(old) - } if third != nil && third != old { // A set from a drain that was never closed. It is not in use and nothing // routes to it any more. s.retire(third, next) } - s.ensureStarted(next) + startErr := s.ensureStarted(next) if s.swaps != nil { s.swaps.Add(ctx, 1, metric.WithAttributes(attribute.String("to", string(next.Name)))) } @@ -399,16 +436,20 @@ func (s *Swapper) swapTo(ctx context.Context, row cluster.CarrierRow) error { } xlog.Info("This replica now uses another carrier; the previous one stays attached while it drains", "carrier", next.Name, "previous", old.Name, "epoch", row.Epoch, "since_prepare", took.Round(time.Millisecond)) + if startErr != nil { + return fmt.Errorf("the carrier is in use but its work did not start: %w", startErr) + } return nil } -// ensureStarted starts the work of a set once. -func (s *Swapper) ensureStarted(set *Set) { +// ensureStarted starts the work of a set once. It returns why the work could not +// start, and tries again at the next call. +func (s *Swapper) ensureStarted(set *Set) error { if set == nil || set.Start == nil { - return + return nil } if _, ok := s.members[set]; ok { - return + return nil } ctx, cancel := context.WithCancelCause(s.base) stop, err := set.Start(ctx) @@ -416,20 +457,28 @@ func (s *Swapper) ensureStarted(set *Set) { cancel(err) // Tried again at the next poll. xlog.Error("Could not start the work of the carrier", "carrier", set.Name, "error", err) - return + s.startErr[set] = err + return err } + delete(s.startErr, set) s.members[set] = &member{cancel: cancel, stop: stop} + return nil } // settle acts on a stable row: it drops a set built for a change that was -// aborted, and releases the previous set once the leader has ended the drain. -func (s *Swapper) settle(row cluster.CarrierRow) { +// aborted, attaches again the carrier that is still draining when this replica +// started during the drain, and releases the previous set once the leader has +// ended the drain. +func (s *Swapper) settle(ctx context.Context, row cluster.CarrierRow) { if s.prepared != nil { if s.prepared != s.draining { closeSet(s.dropListening(s.prepared)) } s.prepared = nil } + if row.Draining != "" && s.draining == nil && row.Draining != s.o.Cur.Load().Name { + s.reattachDraining(ctx, row) + } if row.Draining == "" && s.draining != nil { old := s.draining s.draining = nil @@ -440,6 +489,29 @@ func (s *Swapper) settle(row cluster.CarrierRow) { } } +// reattachDraining builds the carrier that the row says is draining and opens the +// window for it. A replica that starts or restarts during a drain begins on the +// active carrier only, and a worker that is attached to the draining carrier +// only would not be reached until the drain ended. A failure is logged and tried +// again at the next poll; it does not make the replica not ready, because the +// active carrier works. +func (s *Swapper) reattachDraining(ctx context.Context, row cluster.CarrierRow) { + set, err := s.attach(ctx, row, row.Draining) + if err != nil { + xlog.Warn("This replica cannot attach the carrier that is draining; workers that are attached only to it are not reached until the drain ends", + "carrier", row.Draining, "epoch", row.Epoch, "error", err) + return + } + s.draining = set + if s.o.Window != nil { + s.o.Window.Open(set) + } + if err := s.ensureStarted(set); err != nil { + xlog.Warn("The work of the carrier that is draining did not start on this replica", "carrier", row.Draining, "error", err) + } + xlog.Info("This replica attached the carrier that is draining", "carrier", row.Draining, "epoch", row.Epoch) +} + // dropListening detaches every subscription from set. func (s *Swapper) dropListening(set *Set) *Set { if err := s.o.Bus.Release(set); err != nil { @@ -455,6 +527,7 @@ func (s *Swapper) retire(old, next *Set) { s.dropListening(old) m := s.members[old] delete(s.members, old) + delete(s.startErr, old) s.wg.Add(1) go func() { defer s.wg.Done() @@ -462,7 +535,7 @@ func (s *Swapper) retire(old, next *Set) { // A run that is still going loses its carrier here. It is told why, so // that it is not offered to the new carrier to run again. m.cancel(messaging.ErrCarrierReleased) - m.stop() + s.stopBounded(old.Name, m.stop) } s.handoff(old, next) closeSet(old) @@ -482,11 +555,32 @@ func (s *Swapper) retire(old, next *Set) { }() } +// stopBounded waits for stop, and gives up after StopGrace. Work that does not +// stop would hold the hand-over of the queue, the close of the set and the +// release of its connections for as long as it runs. +func (s *Swapper) stopBounded(name cluster.Carrier, stop func()) { + done := make(chan struct{}) + go func() { + defer close(done) + stop() + }() + timer := time.NewTimer(s.o.StopGrace) + defer timer.Stop() + select { + case <-done: + case <-timer.C: + xlog.Warn("The work of a carrier did not stop in time; going on without waiting for it", + "carrier", name, "grace", s.o.StopGrace) + } +} + func (s *Swapper) handoff(from, to *Set) { if from.Handoff == nil || to == nil { return } - ctx, cancel := context.WithTimeout(context.WithoutCancel(s.base), 2*time.Minute) + // Not bound to the context of the replica, so that a stop of the replica does + // not drop the hand-over of the queue; Close bounds it with CloseGrace. + ctx, cancel := context.WithTimeout(s.life, 2*time.Minute) defer cancel() if err := from.Handoff(ctx, to); err != nil { xlog.Warn("Could not hand over what the released carrier held", "from", from.Name, "to", to.Name, "error", err) @@ -520,16 +614,35 @@ func (s *Swapper) Close() { s.o.Window.Close() } - for _, m := range members { + var stopping sync.WaitGroup + for set, m := range members { m.cancel(context.Canceled) - m.stop() + stopping.Go(func() { s.stopBounded(set.Name, m.stop) }) } + stopping.Wait() for _, set := range extra { s.dropListening(set) closeSet(set) } s.doneOnce.Do(func() { close(s.done) }) - s.wg.Wait() + // The hand-over that runs in the background gets CloseGrace to finish. After + // that it is cancelled, and Close returns when it has seen that. + finished := make(chan struct{}) + go func() { s.wg.Wait(); close(finished) }() + timer := time.NewTimer(s.o.CloseGrace) + defer timer.Stop() + select { + case <-finished: + case <-timer.C: + xlog.Warn("The hand-over of a released carrier did not finish before the replica stopped; cancelling it", "grace", s.o.CloseGrace) + s.lifeCancel() + select { + case <-finished: + case <-time.After(s.o.CloseGrace): + xlog.Warn("The hand-over of a released carrier ignored its cancel; the replica stops without it") + } + } + s.lifeCancel() } func closeSet(set *Set) { diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go index 4152af932288..d6dd95cc1b39 100644 --- a/core/services/carrier/swap_test.go +++ b/core/services/carrier/swap_test.go @@ -38,6 +38,51 @@ func newFakeNet() *fakeNet { }} } +// errUsedAfterClose is what a fake carrier answers once its set is closed. A real +// carrier answers with an error or a panic then, and a spec about "never reaches +// a closed set" must be able to see that. +var errUsedAfterClose = errors.New("the carrier was used after it was closed") + +type closedAwareBroadcaster struct { + messaging.Broadcaster + closed *atomic.Bool +} + +func (b closedAwareBroadcaster) Publish(subject string, data any) error { + if b.closed.Load() { + return errUsedAfterClose + } + return b.Broadcaster.Publish(subject, data) +} + +type closedAwareQueue struct { + messaging.WorkQueue + closed *atomic.Bool +} + +func (q closedAwareQueue) Enqueue(ctx context.Context, kind messaging.WorkKind, payload any) error { + if q.closed.Load() { + return errUsedAfterClose + } + return q.WorkQueue.Enqueue(ctx, kind, payload) +} + +// windowSpy is the window of a replica. It notes whether the pointer to the set in +// use still named the draining set when the window opened. +type windowSpy struct { + *carrier.Window + n *replicaNode +} + +func (w windowSpy) Open(prev *carrier.Set) { + if w.n.cur.Load() == prev { + w.n.windowRec.rec("window-open:before-the-pointer-moved") + } else { + w.n.windowRec.rec("window-open:after-the-pointer-moved") + } + w.Window.Open(prev) +} + // life records what the replica did to the sets it built. type life struct { recorder @@ -63,6 +108,14 @@ type replicaNode struct { handedTo []cluster.Carrier attached map[string]carrier.Attachment reconnect atomic.Int32 + + // hangStop makes the stop of the work of a set wait for releaseStop. + windowRec recorder + hangStop atomic.Bool + releaseStop chan struct{} + // hangHandoff makes the hand-over wait for the end of its context. + hangHandoff atomic.Bool + startFail map[cluster.Carrier]error } func (n *replicaNode) build(_ context.Context, row cluster.CarrierRow, target cluster.Carrier) (*carrier.Set, error) { @@ -84,23 +137,39 @@ func (n *replicaNode) build(_ context.Context, row cluster.CarrierRow, target cl func (n *replicaNode) newSet(name cluster.Carrier, epoch int64) *carrier.Set { f := newFakeCarrier(name, epoch) bus := n.net.bus[name] - f.set.Broadcaster = bus + closed := &atomic.Bool{} + f.commands.closed = closed + f.set.Broadcaster = closedAwareBroadcaster{Broadcaster: bus, closed: closed} + f.set.WorkQueue = closedAwareQueue{WorkQueue: f.queue, closed: closed} f.set.OnReconnect = bus.OnReconnect f.set.Start = func(ctx context.Context) (func(), error) { + n.mu.Lock() + failure := n.startFail[name] + n.mu.Unlock() + if failure != nil { + return nil, failure + } n.life.rec("start:" + string(name)) return func() { + if n.hangStop.Load() { + <-n.releaseStop + } released := errors.Is(context.Cause(ctx), messaging.ErrCarrierReleased) n.life.rec(fmt.Sprintf("stop:%s:released=%t", name, released)) }, nil } - f.set.Handoff = func(_ context.Context, next *carrier.Set) error { + f.set.Handoff = func(hctx context.Context, next *carrier.Set) error { n.mu.Lock() n.handedTo = append(n.handedTo, next.Name) n.mu.Unlock() n.life.rec("handoff:" + string(name) + "->" + string(next.Name)) + if n.hangHandoff.Load() { + <-hctx.Done() + n.life.rec("handoff-cancelled:" + string(name)) + } return nil } - f.set.Close = func() { n.life.rec("close:" + string(name)) } + f.set.Close = func() { closed.Store(true); n.life.rec("close:" + string(name)) } n.sets[name] = append(n.sets[name], f) return f.set } @@ -138,6 +207,8 @@ var _ = Describe("The swap of a replica", func() { net *fakeNet nodes []*replicaNode timings cluster.Timings + // stopGrace and closeGrace are the bounds of the swappers that join starts. + stopGrace, closeGrace time.Duration ) // join starts a replica on the carrier that the row names, as initDistributed does. @@ -147,6 +218,7 @@ var _ = Describe("The swap of a replica", func() { id: id, net: net, life: &life{}, builds: map[cluster.Carrier]int{}, buildErr: map[cluster.Carrier]error{}, sets: map[cluster.Carrier][]*fakeCarrier{}, attached: map[string]carrier.Attachment{}, + releaseStop: make(chan struct{}), startFail: map[cluster.Carrier]error{}, } row, err := store.Get(ctx) Expect(err).ToNot(HaveOccurred()) @@ -164,9 +236,10 @@ var _ = Describe("The swap of a replica", func() { n.cmds.UseWindow(n.window) n.queue = carrier.NewWorkQueue(&n.cur) n.swapper, err = carrier.NewSwapper(carrier.SwapperOptions{ - Cur: &n.cur, Bus: n.bus, Window: n.window, Rows: store, + Cur: &n.cur, Bus: n.bus, Window: windowSpy{n.window, n}, Rows: store, Ready: readyReporter{reg: reg, id: id}, Build: n.build, Interval: time.Hour, Settle: 50 * time.Millisecond, + StopGrace: stopGrace, CloseGrace: closeGrace, }) Expect(err).ToNot(HaveOccurred()) n.bus.OnReconnect(func() { n.reconnect.Add(1) }) @@ -236,6 +309,7 @@ var _ = Describe("The swap of a replica", func() { reg = cluster.NewRegistry(db) net = newFakeNet() nodes = nil + stopGrace, closeGrace = 0, 0 timings = cluster.Timings{PrepareTimeout: time.Minute, TransitionWindow: time.Minute, MaxDrain: time.Minute} sw, err = cluster.NewSwitch(cluster.SwitchOptions{Store: store, Registry: reg, Timings: func() cluster.Timings { return timings }}) Expect(err).ToNot(HaveOccurred()) @@ -243,6 +317,11 @@ var _ = Describe("The swap of a replica", func() { AfterEach(func() { for _, n := range nodes { + select { + case <-n.releaseStop: + default: + close(n.releaseStop) + } n.swapper.Close() } }) @@ -627,14 +706,120 @@ var _ = Describe("The swap of a replica", func() { }) }) + Describe("the robustness of a swap", func() { + It("goes on to hand over and to close a set whose work does not stop", func() { + stopGrace = 200 * time.Millisecond + a := join("a") + a.hangStop.Store(true) + settle(cluster.CarrierTunnel) + expire() + + Eventually(a.life.seen, "10s").Should(ContainElements("handoff:nats->tunnel", "close:nats"), + "the stop of the work hangs, and the hand-over and the close do not wait for it") + }) + + It("returns from Close in a bounded time when the hand-over does not end", func() { + closeGrace = 300 * time.Millisecond + a := join("a") + a.hangHandoff.Store(true) + settle(cluster.CarrierTunnel) + expire() + Eventually(a.life.seen, "5s").Should(ContainElement("handoff:nats->tunnel")) + + closed := make(chan struct{}) + go func() { defer close(closed); a.swapper.Close() }() + Eventually(closed, "5s").Should(BeClosed()) + Eventually(a.life.seen, "5s").Should(ContainElement("handoff-cancelled:nats"), "the hand-over is told to stop") + }) + + It("opens the window before the pointer moves, so that no call meets a gap", func() { + a := join("a") + settle(cluster.CarrierTunnel) + Expect(a.windowRec.seen()).To(ContainElement("window-open:before-the-pointer-moved")) + Expect(a.windowRec.seen()).ToNot(ContainElement("window-open:after-the-pointer-moved")) + }) + + It("lets go of the reconnect hook of a set that was released", func() { + a := join("a") + base := a.bus.HookedCount() + settle(cluster.CarrierTunnel) + Expect(a.bus.HookedCount()).To(Equal(base+1), "the tunnel set is hooked, the NATS set is draining") + expire() + Eventually(a.bus.HookedCount, "5s").Should(Equal(base), "the released set is not kept") + }) + + It("lets go of the hook of a set that was built for a change that was aborted", func() { + a := join("a") + base := a.bus.HookedCount() + request(cluster.CarrierTunnel) + poll() + Expect(a.bus.HookedCount()).To(Equal(base + 1)) + _, err := sw.Abort(ctx, "admin") + Expect(err).ToNot(HaveOccurred()) + poll() + Expect(a.bus.HookedCount()).To(Equal(base)) + }) + + It("attaches the draining carrier again when a replica starts during a drain, and opens the window", func() { + a := join("a") + settle(cluster.CarrierTunnel) + Expect(row().Draining).To(Equal(cluster.CarrierNATS)) + + // A replica that starts now begins on the active carrier only. + b := join("b") + b.attached["on-nats"] = carrier.Attachment{NATS: true} + Expect(b.window.Previous()).To(BeNil()) + poll(b) + + Expect(b.buildCount(cluster.CarrierNATS)).To(Equal(1)) + Expect(b.window.Previous()).ToNot(BeNil()) + Expect(b.cmds.PingNode("on-nats")).To(Succeed()) + Expect(b.fake(cluster.CarrierNATS).commands.seen()).To(Equal([]string{"PingNode"}), "a worker that is only on the draining carrier is reached") + + // The drain ends: b lets the set go like a. + expire() + Eventually(b.life.seen, "5s").Should(ContainElement("close:nats")) + Expect(b.window.Previous()).To(BeNil()) + _ = a + }) + + It("says it is not ready when the work of the carrier cannot start, and is ready again when it can", func() { + a := join("a") + a.mu.Lock() + a.startFail[cluster.CarrierTunnel] = errors.New("the claim loop could not start") + a.mu.Unlock() + request(cluster.CarrierTunnel) + poll() + drive() // commit + poll() + + in, err := reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyReason).To(ContainSubstring("claim loop could not start")) + + a.mu.Lock() + a.startFail[cluster.CarrierTunnel] = nil + a.mu.Unlock() + poll() + in, err = reg.Get(ctx, "a") + Expect(err).ToNot(HaveOccurred()) + Expect(in.ReadyReason).To(BeEmpty()) + Expect(a.life.seen()).To(ContainElement("start:tunnel")) + }) + }) + Describe("a call that runs while the set changes", func() { It("never blocks and never reaches a closed set", func() { a, b := join("a"), join("b") + // A worker that is on the carrier that drains only: its calls go to the + // draining set through the window, until the drain ends. + a.attached["on-nats"] = carrier.Attachment{NATS: true} stop := make(chan struct{}) var wg sync.WaitGroup var calls atomic.Int64 for range 4 { wg.Go(func() { + defer GinkgoRecover() for { select { case <-stop: @@ -643,11 +828,17 @@ var _ = Describe("The swap of a replica", func() { } Expect(a.bus.Publish(swapSubject, "x")).To(Succeed()) Expect(a.queue.Enqueue(ctx, messaging.WorkMCPCI, "p")).To(Succeed()) + // The fake carriers answer a call to a closed set with an error, as + // a real one does, so a call that reaches one fails this spec. + Expect(a.cmds.PingNode("on-nats")).To(Succeed()) calls.Add(1) } }) } settle(cluster.CarrierTunnel) + time.Sleep(100 * time.Millisecond) + a.attached["on-nats"] = carrier.Attachment{Tunnel: true} + expire() settle(cluster.CarrierNATS) close(stop) wg.Wait() From 2367992d2c110235356fd75c286dc8d833444ea7 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:20:03 +0000 Subject: [PATCH 140/162] refactor(carrier): remove Fanout.Ready, which nothing reads The ready report of a replica never consumed the check, and only its own spec called it. The LISTEN connection is checked when the carrier is built, and a lost connection is redialled by the bus. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/pgbus.go | 18 ++++-------------- core/services/carrier/pgbus_test.go | 9 --------- 2 files changed, 4 insertions(+), 23 deletions(-) diff --git a/core/services/carrier/pgbus.go b/core/services/carrier/pgbus.go index 4b48b5b1a74f..9fa89236a9a2 100644 --- a/core/services/carrier/pgbus.go +++ b/core/services/carrier/pgbus.go @@ -26,18 +26,14 @@ type PgbusOptions struct { } // Fanout is the fan-out member of a Set for a carrier that has no NATS server: -// the broadcaster, the reconnect hook the holder reads, a ready check and the -// close. The other members of the set come from the other parts of the carrier. +// the broadcaster, the reconnect hook the holder reads and the close. The other +// members of the set come from the other parts of the carrier. type Fanout struct { // Broadcaster delivers through LISTEN and NOTIFY. Broadcaster messaging.Broadcaster // OnReconnect registers a callback that runs after the LISTEN connection was // found again. It is the OnReconnect of the Set. OnReconnect func(func()) - // Ready reports an error when the LISTEN connection is down. The report that - // a replica sends for a change of carrier reads it. A caller must not decide - // on it anything about a worker. - Ready func() error // Dropped reports how many broadcasts this replica received and lost. Dropped func() uint64 // Close ends the LISTEN connection. It is the Close of the Set. @@ -79,13 +75,7 @@ func NewPgbusFanout(ctx context.Context, o PgbusOptions) (*Fanout, error) { return &Fanout{ Broadcaster: bus, OnReconnect: bus.OnReconnect, - Ready: func() error { - if !bus.IsConnected() { - return errors.New("the LISTEN connection to the database is down") - } - return nil - }, - Dropped: bus.Dropped, - Close: bus.Close, + Dropped: bus.Dropped, + Close: bus.Close, }, nil } diff --git a/core/services/carrier/pgbus_test.go b/core/services/carrier/pgbus_test.go index 9b2bee888fd1..444772656e98 100644 --- a/core/services/carrier/pgbus_test.go +++ b/core/services/carrier/pgbus_test.go @@ -63,15 +63,6 @@ var _ = Describe("The pgbus fan-out of the tunnel carrier", func() { Eventually(func() int64 { return listeners(db) }, 10*time.Second).Should(BeZero()) }) - It("reports whether the LISTEN connection is up, for the ready report", func() { - fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) - Expect(err).ToNot(HaveOccurred()) - Expect(fan.Ready()).To(Succeed()) - - fan.Close() - Expect(fan.Ready()).To(HaveOccurred()) - }) - It("creates the spill table when it is built", func() { fan, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: db, DSN: dsn}) Expect(err).ToNot(HaveOccurred()) From 7e3cae2ef3314a6ff0e58ea94a4880b10510ce21 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:21:14 +0000 Subject: [PATCH 141/162] fix(cluster): refuse a wait that is zero, negative or too short, before storing anything The endpoint of the cluster settings checked that a wait parsed, and the store refused one that was not positive only after it had stored the settings before it. A request with a valid address and a wait of 0s stored the address, answered 500, and left half of the request in place. A value such as 1ns passed both checks, and it aborted every later change before a replica could see it. The endpoint now runs the same check as the store on every part of the request before it writes. A wait must be at least five seconds. The timings that a replica reads ignore a stored value below that. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../http/endpoints/localai/cluster_carrier.go | 22 +++++++---------- .../endpoints/localai/cluster_carrier_test.go | 13 ++++++++++ core/services/cluster/natsurl_test.go | 24 +++++++++++++++++++ core/services/cluster/settings.go | 16 ++++++++++--- 4 files changed, 59 insertions(+), 16 deletions(-) diff --git a/core/http/endpoints/localai/cluster_carrier.go b/core/http/endpoints/localai/cluster_carrier.go index d866ec50d74b..ab75d4127bf9 100644 --- a/core/http/endpoints/localai/cluster_carrier.go +++ b/core/http/endpoints/localai/cluster_carrier.go @@ -243,26 +243,22 @@ func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, p type change struct { key string val *string - url bool } changes := []change{ - {cluster.SettingNATSURL, req.NATSURL, true}, - {cluster.SettingNATSWorkerURL, req.NATSWorkerURL, true}, - {cluster.SettingPrepareTimeout, req.PrepareTimeout, false}, - {cluster.SettingTransitionWindow, req.TransitionWindow, false}, - {cluster.SettingMaxDrain, req.MaxDrain, false}, + {cluster.SettingNATSURL, req.NATSURL}, + {cluster.SettingNATSWorkerURL, req.NATSWorkerURL}, + {cluster.SettingPrepareTimeout, req.PrepareTimeout}, + {cluster.SettingTransitionWindow, req.TransitionWindow}, + {cluster.SettingMaxDrain, req.MaxDrain}, } - // Validate everything before storing anything. + // Validate everything before storing anything. The store checks again, and a + // request that it would refuse half way must not be accepted here. for _, ch := range changes { if ch.val == nil || *ch.val == "" { continue } - if ch.url { - if err := cluster.CheckNATSURL(*ch.val); err != nil { - return carrierError(c, http.StatusBadRequest, fmt.Sprintf("%s: %v", ch.key, err), nil) - } - } else if _, err := time.ParseDuration(*ch.val); err != nil { - return carrierError(c, http.StatusBadRequest, fmt.Sprintf("%s: %v", ch.key, err), nil) + if err := cluster.CheckSetting(ch.key, *ch.val); err != nil { + return carrierError(c, http.StatusBadRequest, err.Error(), nil) } } diff --git a/core/http/endpoints/localai/cluster_carrier_test.go b/core/http/endpoints/localai/cluster_carrier_test.go index 3518ed2643ad..a514977577b2 100644 --- a/core/http/endpoints/localai/cluster_carrier_test.go +++ b/core/http/endpoints/localai/cluster_carrier_test.go @@ -302,6 +302,19 @@ var _ = Describe("The admin API of the carrier", func() { Expect(code).To(Equal(http.StatusBadRequest)) }) + It("refuses a wait that is not positive or is too short, with a 400, and stores none of the request", func() { + for _, bad := range []string{"0s", "-5s", "1ns", "1ms", "4s"} { + body := `{"nats_url":"nats://broker2:4222","max_drain":"20m","prepare_timeout":"` + bad + `"}` + code, resp := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, body) + Expect(code).To(Equal(http.StatusBadRequest), bad) + Expect(resp["error"]).To(ContainSubstring("prepare_timeout"), bad) + Expect(settings.m).ToNot(HaveKey(cluster.SettingMaxDrain), "%s: nothing is stored when one part is refused", bad) + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("broker:4222"), bad) + } + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"prepare_timeout":"5s"}`) + Expect(code).To(Equal(http.StatusOK)) + }) + It("reads the settings with the password of the URL masked", func() { code, body := call(GetClusterSettingsEndpoint(settings), http.MethodGet, "") Expect(code).To(Equal(http.StatusOK)) diff --git a/core/services/cluster/natsurl_test.go b/core/services/cluster/natsurl_test.go index a47094063b3b..36dc9a689eae 100644 --- a/core/services/cluster/natsurl_test.go +++ b/core/services/cluster/natsurl_test.go @@ -2,6 +2,7 @@ package cluster_test import ( "context" + "time" "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/testutil" @@ -48,3 +49,26 @@ var _ = Describe("the NATS address of the cluster", func() { } }) }) + +var _ = Describe("the waits of a change in the cluster settings", func() { + It("refuses a wait that is not a sensible duration", func() { + store, err := cluster.NewSettingsStore(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + ctx := context.Background() + for _, bad := range []string{"0s", "-1m", "1ns", "999ms", "soon"} { + Expect(store.Set(ctx, cluster.SettingMaxDrain, bad, "x")).ToNot(Succeed(), bad) + Expect(cluster.CheckSetting(cluster.SettingMaxDrain, bad)).ToNot(Succeed(), bad) + } + Expect(store.Set(ctx, cluster.SettingMaxDrain, "5s", "x")).To(Succeed()) + }) + + It("ignores a stored wait below the minimum when it reads the timings", func() { + db := testutil.SetupTestDB() + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + Expect(db.Create(&cluster.Setting{Key: cluster.SettingPrepareTimeout, Value: "1ns"}).Error).To(Succeed()) + t, err := store.Timings(context.Background(), cluster.Timings{PrepareTimeout: 40 * time.Second}) + Expect(err).ToNot(HaveOccurred()) + Expect(t.PrepareTimeout).To(Equal(40 * time.Second)) + }) +}) diff --git a/core/services/cluster/settings.go b/core/services/cluster/settings.go index 1be5b681c28c..cd0cb2cd0bad 100644 --- a/core/services/cluster/settings.go +++ b/core/services/cluster/settings.go @@ -44,17 +44,27 @@ var knownSettings = map[string]func(string) error{ SettingMaxDrain: durationSetting, } +// MinTiming is the shortest wait of a change that the cluster accepts. The +// replicas read the row every two seconds, so a shorter wait would abort every +// change before a replica could see it, and a value such as 1ns is a typo for +// a value that is meant to be long. +const MinTiming = 5 * time.Second + func durationSetting(v string) error { d, err := time.ParseDuration(v) if err != nil { return fmt.Errorf("not a duration: %w", err) } - if d <= 0 { - return errors.New("a duration must be positive") + if d < MinTiming { + return fmt.Errorf("a wait must be at least %s, not %s", MinTiming, d) } return nil } +// CheckSetting says whether value may be stored for key. It stores nothing, so a +// caller with several settings can refuse all of them before it writes one. +func CheckSetting(key, value string) error { return checkSetting(key, value) } + // Timings are the waits of a change. type Timings struct { // PrepareTimeout is how long prepare waits for every live replica. @@ -202,7 +212,7 @@ func (s *SettingsStore) Timings(ctx context.Context, fallback Timings) (Timings, SettingMaxDrain: &t.MaxDrain, } { if v, ok := all[key]; ok { - if d, err := time.ParseDuration(v); err == nil && d > 0 { + if d, err := time.ParseDuration(v); err == nil && d >= MinTiming { *dst = d } } From de1d1c58095c50c5bffb0596667155979c81461e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:21:37 +0000 Subject: [PATCH 142/162] fix(cluster): keep the note of a change when its drain ends The end of a drain replaced the note of the change with the word "drained", so an admin who looked at the status after a forced change no longer saw that it was forced or which replica it left behind. The drain ending is recorded by the log and the counter. The note stays. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/cluster/switch.go | 4 +++- core/services/cluster/switch_test.go | 24 ++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index 5bc43a416fc2..4f166e2315df 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -578,8 +578,10 @@ func (s *Switch) driveDrain(ctx context.Context, row CarrierRow, now time.Time) if row.DrainingUntil != nil && now.Before(*row.DrainingUntil) { return nil } + // The note of the change stays: it says how that change ended, and the drain + // ending is an ordinary event that the log and the counter record. _, err := s.transition(ctx, row.Epoch, Change{ - Active: row.Active, State: StateStable, Note: "drained", By: row.ChangedBy, + Active: row.Active, State: StateStable, Note: row.Note, By: row.ChangedBy, }) if err == nil { s.count("drained") diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go index 06587cf3f704..84d31346ca7d 100644 --- a/core/services/cluster/switch_test.go +++ b/core/services/cluster/switch_test.go @@ -541,6 +541,30 @@ var _ = Describe("The switch of the carrier", func() { Expect(got.Note).To(ContainSubstring("b")) }) + It("keeps the note of the change when the drain ends", func() { + timings.PrepareTimeout = 300 * time.Millisecond + timings.MaxDrain = 500 * time.Millisecond + replica("a", "", "") + replica("b", "", "") + prepare(true) + ready("a", row().Epoch, "") + Eventually(func() cluster.State { + Expect(sw.Drive(ctx)).To(Succeed()) + return row().State + }, "10s", "100ms").Should(Equal(cluster.StateCommit)) + noted := row().Note + Expect(noted).To(ContainSubstring("without")) + ready("a", row().Epoch, "") + ready("b", row().Epoch, "") + Expect(sw.Drive(ctx)).To(Succeed()) + + Eventually(func() cluster.Carrier { + Expect(sw.Drive(ctx)).To(Succeed()) + return row().Draining + }, "10s", "100ms").Should(BeEmpty()) + Expect(row().Note).To(Equal(noted), "the note says how the change ended, and the end of the drain does not replace it") + }) + It("does not wait for the timeout for a forced change when a replica says it cannot build the target", func() { replica("a", "", "") replica("b", "", "") From 7c3640e2addd5364df1bb57667d01adb3788c2cc Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:22:34 +0000 Subject: [PATCH 143/162] fix(nodes): read the age of a worker heartbeat on the clock of the database The switch of the carrier counted a worker as alive when its last heartbeat was younger than the stale limit by the clock of the replica that read it. The heartbeat is stamped by whichever replica handled it, so a skew between two replicas made a live worker look stale, or a dead one look alive. The age is now read on the clock of the database, which every replica shares, and falls back to the local clock only when the database does not answer. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/follow_report.go | 23 ++++++++++++++++++++--- core/services/nodes/follow_report_test.go | 23 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/core/services/nodes/follow_report.go b/core/services/nodes/follow_report.go index 3230f6550900..09f76bcd1e11 100644 --- a/core/services/nodes/follow_report.go +++ b/core/services/nodes/follow_report.go @@ -99,12 +99,27 @@ func (r *NodeRegistry) SetCarrierReport(ctx context.Context, nodeID string, rep type SwitchWorkers struct { reg *NodeRegistry stale time.Duration + // now is the clock that the age of a heartbeat is read on. It is the clock of + // the database, which every replica shares. The clock of the replica that reads + // is not that of the replica that wrote the heartbeat, and a skew between the + // two would turn a live worker into a stale one, or the other way round. + now func(ctx context.Context) time.Time } // NewSwitchWorkers returns the reader. A worker whose heartbeat is older than // stale is not alive for this purpose. func NewSwitchWorkers(reg *NodeRegistry, stale time.Duration) *SwitchWorkers { - return &SwitchWorkers{reg: reg, stale: stale} + return &SwitchWorkers{reg: reg, stale: stale, now: reg.databaseTime} +} + +// databaseTime returns the time of the database. If the database does not answer it +// returns the local time, because a change of carrier must not fail on it. +func (r *NodeRegistry) databaseTime(ctx context.Context) time.Time { + var now time.Time + if err := r.db.WithContext(ctx).Raw("SELECT now()").Scan(&now).Error; err != nil || now.IsZero() { + return time.Now() + } + return now } var _ cluster.WorkerSource = (*SwitchWorkers)(nil) @@ -125,10 +140,11 @@ func (s *SwitchWorkers) Workers(ctx context.Context) ([]cluster.WorkerInfo, erro if err != nil { return nil, err } + now := s.now(ctx) var out []cluster.WorkerInfo for i := range list { n := &list[i] - if n.Status == StatusPending || n.Status == StatusOffline || time.Since(n.LastHeartbeat) > s.stale { + if n.Status == StatusPending || n.Status == StatusOffline || now.Sub(n.LastHeartbeat) > s.stale { continue } attached := s.attached(n) @@ -156,9 +172,10 @@ func (s *SwitchWorkers) AgentsAttached(ctx context.Context, c cluster.Carrier) ( if err != nil { return false, err } + now := s.now(ctx) for i := range list { n := &list[i] - if n.NodeType != NodeTypeAgent || n.Status == StatusPending || n.Status == StatusOffline || time.Since(n.LastHeartbeat) > s.stale { + if n.NodeType != NodeTypeAgent || n.Status == StatusPending || n.Status == StatusOffline || now.Sub(n.LastHeartbeat) > s.stale { continue } for _, a := range s.attached(n) { diff --git a/core/services/nodes/follow_report_test.go b/core/services/nodes/follow_report_test.go index 04775033b43d..441ec92561d0 100644 --- a/core/services/nodes/follow_report_test.go +++ b/core/services/nodes/follow_report_test.go @@ -60,6 +60,29 @@ var _ = Describe("The workers as the switch of the carrier sees them", func() { Expect(byName(list)).To(HaveKey("live")) }) + It("reads the age of a heartbeat on the clock of the database and not on the clock of this replica", func() { + n := register("agent-1", NodeTypeAgent, "10.0.0.9:50051") + beat := time.Now() + Expect(db.Model(&BackendNode{}).Where("id = ?", n.ID).Update("last_heartbeat", beat).Error).To(Succeed()) + + // This replica's clock says the heartbeat is fresh. The clock of the database + // is ten minutes ahead, which makes it stale. + skewed := NewSwitchWorkers(reg, 5*time.Minute) + skewed.now = func(context.Context) time.Time { return beat.Add(10 * time.Minute) } + list, err := skewed.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(list).To(BeEmpty()) + on, err := skewed.AgentsAttached(ctx, cluster.CarrierNATS) + Expect(err).ToNot(HaveOccurred()) + Expect(on).To(BeFalse()) + + // And the other way: a clock that is ahead here does not hide a live worker. + real := NewSwitchWorkers(reg, 5*time.Minute) + list, err = real.Workers(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(list).To(HaveLen(1)) + }) + It("counts a worker that reports no capabilities as one that cannot follow, and as one that is on NATS", func() { register("old", NodeTypeBackend, "10.0.0.1:50051") list, err := workers.Workers(ctx) From e514139cb1c347bb44ac5aabd12bd522346fd5ef Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:23:18 +0000 Subject: [PATCH 144/162] fix(nodes): convert the replica index to int32 with a bound gosec flagged the plain conversion of the replica index in four requests (G115). The index is a small count, so the conversion now clamps a value that does not fit and does not wrap it, which could have named another replica. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/nodes/follow_report_test.go | 10 ++++++++++ core/services/nodes/unloader.go | 22 ++++++++++++++++++---- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/core/services/nodes/follow_report_test.go b/core/services/nodes/follow_report_test.go index 441ec92561d0..6b3f1b638f24 100644 --- a/core/services/nodes/follow_report_test.go +++ b/core/services/nodes/follow_report_test.go @@ -3,6 +3,7 @@ package nodes import ( "context" "errors" + "math" "time" . "github.com/onsi/ginkgo/v2" @@ -227,3 +228,12 @@ var _ = Describe("The direct client factory for a worker that holds a tunnel", f Expect(asked).To(Equal(1)) }) }) + +var _ = Describe("the index of a replica on the wire", func() { + It("keeps a small index and clamps one that does not fit, instead of wrapping it", func() { + Expect(replicaIndex32(0)).To(Equal(int32(0))) + Expect(replicaIndex32(7)).To(Equal(int32(7))) + Expect(replicaIndex32(math.MaxInt32 + 1)).To(Equal(int32(math.MaxInt32))) + Expect(replicaIndex32(math.MinInt32 - 1)).To(Equal(int32(math.MinInt32))) + }) +}) diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index d11b53796b82..c9c3fbd3f2ea 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "math" "time" "github.com/mudler/LocalAI/core/services/galleryop" @@ -204,6 +205,19 @@ func (a *nodeControl) OperationControl(nodeID string, req workerctl.OperationReq const operationControlTimeout = 5 * time.Second +// replicaIndex32 converts the index of a replica for the wire. An index is a +// small count; a value outside the range of the field is clamped to its end and +// not wrapped, so that it can never name another replica. +func replicaIndex32(i int) int32 { + switch { + case i > math.MaxInt32: + return math.MaxInt32 + case i < math.MinInt32: + return math.MinInt32 + } + return int32(i) +} + // InstallBackendOp is InstallBackend for a load operation: the request carries // the operation id and the longest the load may run, so the worker can bound it. func (a *nodeControl) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) { @@ -211,7 +225,7 @@ func (a *nodeControl) InstallBackendOp(nodeID, backendType, modelID, galleriesJS Backend: backendType, ModelID: modelID, BackendGalleries: galleriesJSON, - ReplicaIndex: int32(replicaIndex), + ReplicaIndex: replicaIndex32(replicaIndex), OpID: opID, OperationID: operationID, DeadlineMs: deadline.Milliseconds(), @@ -313,7 +327,7 @@ func (a *nodeControl) InstallBackend( URI: uri, Name: name, Alias: alias, - ReplicaIndex: int32(replicaIndex), + ReplicaIndex: replicaIndex32(replicaIndex), OpID: opID, }, onProgress) } @@ -356,7 +370,7 @@ func (a *nodeControl) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, na URI: uri, Name: name, Alias: alias, - ReplicaIndex: int32(replicaIndex), + ReplicaIndex: replicaIndex32(replicaIndex), OpID: opID, }, a.upgradeTimeout, onProgress) @@ -393,7 +407,7 @@ func (a *nodeControl) InstallBackendForce(nodeID, backendType, galleriesJSON, ur URI: uri, Name: name, Alias: alias, - ReplicaIndex: int32(replicaIndex), + ReplicaIndex: replicaIndex32(replicaIndex), Force: true, OpID: opID, }, a.upgradeTimeout, onProgress) From 7c5f4864417bcdddab863446876dfa2737d58277 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:24:15 +0000 Subject: [PATCH 145/162] fix(cluster): count a replica's report only if it came after the check was asked A dry run or a request asks every replica what it can build, waits up to five seconds, and then reads the reports. The preflight accepted a report up to 30 seconds old, so a replica that never heard the question passed on the strength of an old answer. The endpoint now measures how long the wait took, and the preflight counts a report only if it is younger than that. A replica that did not answer shows as a blocker that can be forced. A caller that asked nothing keeps the old age. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../http/endpoints/localai/cluster_carrier.go | 10 +++++-- .../endpoints/localai/cluster_carrier_test.go | 28 ++++++++++++++++++- core/services/cluster/switch.go | 27 +++++++++++++++--- core/services/cluster/switch_test.go | 26 +++++++++++++++++ 4 files changed, 84 insertions(+), 7 deletions(-) diff --git a/core/http/endpoints/localai/cluster_carrier.go b/core/http/endpoints/localai/cluster_carrier.go index ab75d4127bf9..1a3f2692aa2b 100644 --- a/core/http/endpoints/localai/cluster_carrier.go +++ b/core/http/endpoints/localai/cluster_carrier.go @@ -19,6 +19,9 @@ import ( type CarrierSwitchAdmin interface { Status(ctx context.Context) (cluster.Report, error) Preflight(ctx context.Context, target cluster.Carrier) (cluster.Report, error) + // PreflightWithin is Preflight where a report of what a replica can build + // counts only when it is younger than within. + PreflightWithin(ctx context.Context, target cluster.Carrier, within time.Duration) (cluster.Report, error) Request(ctx context.Context, req cluster.Request) (cluster.CarrierRow, cluster.Report, error) Abort(ctx context.Context, by string) (cluster.CarrierRow, error) } @@ -162,17 +165,20 @@ func SwitchCarrierEndpoint(sw CarrierSwitchAdmin, prober ReplicaProber) echo.Han } // The replicas write what they can build when they are asked, so the answer // of the preflight is about now and not about the last time they looked. + // A report counts as the answer only if it was written after the question. + asked := time.Now() prober.ProbeReplicas(ctx) + answeredWithin := time.Since(asked) if req.DryRun { - report, err := sw.Preflight(ctx, target) + report, err := sw.PreflightWithin(ctx, target, answeredWithin) if err != nil { return carrierError(c, http.StatusInternalServerError, err.Error(), nil) } return c.JSON(http.StatusOK, report) } - row, report, err := sw.Request(ctx, cluster.Request{Target: target, By: by, Force: req.Force}) + row, report, err := sw.Request(ctx, cluster.Request{Target: target, By: by, Force: req.Force, ReportedWithin: answeredWithin}) var blocked *cluster.BlockedError switch { case errors.As(err, &blocked): diff --git a/core/http/endpoints/localai/cluster_carrier_test.go b/core/http/endpoints/localai/cluster_carrier_test.go index a514977577b2..895d64824a1a 100644 --- a/core/http/endpoints/localai/cluster_carrier_test.go +++ b/core/http/endpoints/localai/cluster_carrier_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "strings" "sync" + "time" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/services/cluster" @@ -21,6 +22,7 @@ type fakeSwitch struct { report cluster.Report row cluster.CarrierRow requests []cluster.Request + within []time.Duration aborts []string requestE error abortE error @@ -44,6 +46,13 @@ func (f *fakeSwitch) Preflight(_ context.Context, target cluster.Carrier) (clust return r, nil } +func (f *fakeSwitch) PreflightWithin(ctx context.Context, target cluster.Carrier, within time.Duration) (cluster.Report, error) { + f.mu.Lock() + f.within = append(f.within, within) + f.mu.Unlock() + return f.Preflight(ctx, target) +} + func (f *fakeSwitch) Request(_ context.Context, req cluster.Request) (cluster.CarrierRow, cluster.Report, error) { f.note("request") f.mu.Lock() @@ -59,9 +68,13 @@ func (f *fakeSwitch) Abort(_ context.Context, by string) (cluster.CarrierRow, er return f.row, f.abortE } -type fakeProber struct{ order *[]string } +type fakeProber struct { + order *[]string + delay time.Duration +} func (f fakeProber) ProbeReplicas(context.Context) { + time.Sleep(f.delay) if f.order != nil { *f.order = append(*f.order, "probe") } @@ -178,6 +191,19 @@ var _ = Describe("The admin API of the carrier", func() { Expect(sw.requests).To(BeEmpty()) }) + It("counts only a report that is younger than the check it just asked for", func() { + slow := fakeProber{delay: 40 * time.Millisecond} + code, _ := call(SwitchCarrierEndpoint(sw, slow), http.MethodPost, `{"target":"tunnel","dry_run":true}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(sw.within).To(HaveLen(1)) + Expect(sw.within[0]).To(BeNumerically(">=", 40*time.Millisecond), "the age that counts starts at the question") + Expect(sw.within[0]).To(BeNumerically("<", 5*time.Second)) + + code, _ = call(SwitchCarrierEndpoint(sw, slow), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusAccepted)) + Expect(sw.requests[0].ReportedWithin).To(BeNumerically(">=", 40*time.Millisecond)) + }) + It("starts the change, and passes force on", func() { code, body := call(SwitchCarrierEndpoint(sw, fakeProber{order: &order}), http.MethodPost, `{"target":"tunnel","force":true}`) Expect(code).To(Equal(http.StatusAccepted)) diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index 4f166e2315df..ca7f82224f01 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -165,6 +165,10 @@ type Request struct { By string // Force goes ahead past a replica or a worker that is not ready. Force bool + // ReportedWithin, when it is not zero, is the age under which a report of + // what a replica can build counts as the answer to the check that the caller + // has just asked for. + ReportedWithin time.Duration } // BlockedError is returned by Switch.Request when the preflight blocks it. The @@ -270,18 +274,26 @@ func (s *Switch) observePhase(phase string, d time.Duration) { // Status reports the cluster as it is, with no target. func (s *Switch) Status(ctx context.Context) (Report, error) { - return s.report(ctx, "") + return s.report(ctx, "", 0) } // Preflight reports what a change to target would meet. It changes nothing. func (s *Switch) Preflight(ctx context.Context, target Carrier) (Report, error) { + return s.PreflightWithin(ctx, target, 0) +} + +// PreflightWithin is Preflight for a caller that has just asked the replicas what +// they can build. A replica counts as having answered only when its report is +// younger than within, so that a report from before the question is not taken for +// an answer to it. Zero keeps the age that the switch accepts anyway. +func (s *Switch) PreflightWithin(ctx context.Context, target Carrier, within time.Duration) (Report, error) { if !target.valid() { return Report{}, fmt.Errorf("%w: %q", ErrInvalidCarrier, target) } - return s.report(ctx, target) + return s.report(ctx, target, within) } -func (s *Switch) report(ctx context.Context, target Carrier) (Report, error) { +func (s *Switch) report(ctx context.Context, target Carrier, within time.Duration) (Report, error) { row, err := s.o.Store.Get(ctx) if err != nil { return Report{}, err @@ -340,12 +352,19 @@ func (s *Switch) report(ctx context.Context, target Carrier) (Report, error) { r.Blockers = append(r.Blockers, Blocker{Kind: BlockerSame, Reason: fmt.Sprintf("%s is the active carrier already", target)}) } if len(r.Blockers) == 0 { + maxAge := s.o.AvailabilityMaxAge + if within > 0 && within < maxAge { + maxAge = within + } for _, in := range live { reason, known := in.AvailabilityFor(target) switch { case !known: r.Blockers = append(r.Blockers, Blocker{Kind: BlockerReplica, ID: in.ID, Forceable: true, Reason: fmt.Sprintf("the replica has not reported whether it can use the %s carrier", target)}) + case within > 0 && in.AvailabilityAge > maxAge: + r.Blockers = append(r.Blockers, Blocker{Kind: BlockerReplica, ID: in.ID, Forceable: true, + Reason: fmt.Sprintf("the replica has not answered the check for the %s carrier: its last report is %s old", target, in.AvailabilityAge.Round(time.Second))}) case in.AvailabilityAge > s.o.AvailabilityMaxAge: r.Blockers = append(r.Blockers, Blocker{Kind: BlockerReplica, ID: in.ID, Forceable: true, Reason: fmt.Sprintf("what the replica reported about the %s carrier is stale (%s old)", target, in.AvailabilityAge.Round(time.Second))}) @@ -371,7 +390,7 @@ func (s *Switch) report(ctx context.Context, target Carrier) (Report, error) { // blocks, moves the row to prepare. A blocker that can be forced does not block a // forced request. The row is unchanged when the request is refused. func (s *Switch) Request(ctx context.Context, req Request) (CarrierRow, Report, error) { - report, err := s.Preflight(ctx, req.Target) + report, err := s.PreflightWithin(ctx, req.Target, req.ReportedWithin) if err != nil { return CarrierRow{}, Report{}, err } diff --git a/core/services/cluster/switch_test.go b/core/services/cluster/switch_test.go index 84d31346ca7d..13d64cad4679 100644 --- a/core/services/cluster/switch_test.go +++ b/core/services/cluster/switch_test.go @@ -143,6 +143,32 @@ var _ = Describe("The switch of the carrier", func() { ))) }) + It("counts a report only if it is younger than the check the caller just asked for", func() { + replica("a", "", "") + time.Sleep(1200 * time.Millisecond) // the report is older than the question + + plain, err := sw.Preflight(ctx, cluster.CarrierTunnel) + Expect(err).ToNot(HaveOccurred()) + Expect(plain.OK).To(BeTrue(), "a report of 1 second is fresh enough for a preflight that asked nothing") + + asked, err := sw.PreflightWithin(ctx, cluster.CarrierTunnel, 500*time.Millisecond) + Expect(err).ToNot(HaveOccurred()) + Expect(asked.OK).To(BeFalse()) + Expect(asked.Blockers).To(HaveLen(1)) + Expect(asked.Blockers[0].ID).To(Equal("a")) + Expect(asked.Blockers[0].Reason).To(ContainSubstring("has not answered the check")) + Expect(asked.Blockers[0].Forceable).To(BeTrue()) + + _, _, err = sw.Request(ctx, cluster.Request{Target: cluster.CarrierTunnel, By: "admin", ReportedWithin: 500 * time.Millisecond}) + var blocked *cluster.BlockedError + Expect(errors.As(err, &blocked)).To(BeTrue()) + + replica("a", "", "") // the replica answers the check + fresh, err := sw.PreflightWithin(ctx, cluster.CarrierTunnel, 500*time.Millisecond) + Expect(err).ToNot(HaveOccurred()) + Expect(fresh.OK).To(BeTrue()) + }) + It("ignores what a replica reported long ago", func() { replica("a", "", "") Expect(db.Exec(`UPDATE instances SET availability_at = now() - interval '1 hour' WHERE id = 'a'`).Error).To(Succeed()) From 89150ac73405e952da3fee635bdc2ea7f13df3a2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:39:45 +0000 Subject: [PATCH 146/162] fix(carrier): keep a released set open for the settle time before it is closed The swapper closed the released set at once after the hand-over, and only then waited for the producers that were slow to flip. A call that had taken the draining set from the window, or the old set from the pointer, a moment before they moved could meet a closed set, and a slow producer's work on the old carrier met a closed queue. Under CPU load a spec that calls through the window during the drain failed in 1 run of 10. The set now stays open for the settle time, the second hand-over runs, and the set is closed after it. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/carrier/swap.go | 11 +++++++---- core/services/carrier/swap_test.go | 8 +++++++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/core/services/carrier/swap.go b/core/services/carrier/swap.go index e4cc74f0fd29..aedbc9ff8522 100644 --- a/core/services/carrier/swap.go +++ b/core/services/carrier/swap.go @@ -538,20 +538,23 @@ func (s *Swapper) retire(old, next *Set) { s.stopBounded(old.Name, m.stop) } s.handoff(old, next) - closeSet(old) // Messages across the flip were not ordered, so the consumers read their // tables once. s.o.Bus.NotifyReconnect() + // The set stays open for the settle time. A call that took the old set from + // the window or from the pointer just before they moved is still using it, + // and a producer that was slow to flip can still put work on it. Closing it + // first would fail those calls. select { case <-time.After(s.o.Settle): case <-s.done: } - // A producer that was slow to flip can still have put work on the old - // carrier. The handoff is a compare-and-set, so a second run finds nothing - // that the first took. + // The handoff is a compare-and-set, so a second run finds nothing that the + // first took. if cur := s.o.Cur.Load(); cur != old { s.handoff(old, cur) } + closeSet(old) }() } diff --git a/core/services/carrier/swap_test.go b/core/services/carrier/swap_test.go index d6dd95cc1b39..1005f4f813b6 100644 --- a/core/services/carrier/swap_test.go +++ b/core/services/carrier/swap_test.go @@ -209,6 +209,8 @@ var _ = Describe("The swap of a replica", func() { timings cluster.Timings // stopGrace and closeGrace are the bounds of the swappers that join starts. stopGrace, closeGrace time.Duration + // settleFor is how long a released set stays open. + settleFor time.Duration ) // join starts a replica on the carrier that the row names, as initDistributed does. @@ -238,7 +240,7 @@ var _ = Describe("The swap of a replica", func() { n.swapper, err = carrier.NewSwapper(carrier.SwapperOptions{ Cur: &n.cur, Bus: n.bus, Window: windowSpy{n.window, n}, Rows: store, Ready: readyReporter{reg: reg, id: id}, Build: n.build, - Interval: time.Hour, Settle: 50 * time.Millisecond, + Interval: time.Hour, Settle: settleFor, StopGrace: stopGrace, CloseGrace: closeGrace, }) Expect(err).ToNot(HaveOccurred()) @@ -310,6 +312,7 @@ var _ = Describe("The swap of a replica", func() { net = newFakeNet() nodes = nil stopGrace, closeGrace = 0, 0 + settleFor = 50 * time.Millisecond timings = cluster.Timings{PrepareTimeout: time.Minute, TransitionWindow: time.Minute, MaxDrain: time.Minute} sw, err = cluster.NewSwitch(cluster.SwitchOptions{Store: store, Registry: reg, Timings: func() cluster.Timings { return timings }}) Expect(err).ToNot(HaveOccurred()) @@ -810,6 +813,9 @@ var _ = Describe("The swap of a replica", func() { Describe("a call that runs while the set changes", func() { It("never blocks and never reaches a closed set", func() { + // A call that took the draining set just before the window closed gets this + // long to finish. It is long against the pause of a busy machine. + settleFor = 2 * time.Second a, b := join("a"), join("b") // A worker that is on the carrier that drains only: its calls go to the // draining set through the window, until the drain ends. From 5a5522d86b36a8682ffd6189925c671f78e78b97 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:52:40 +0000 Subject: [PATCH 147/162] feat(mcp): add a read-only get_cluster_carrier tool The assistant can now read which transport a distributed cluster uses, the state and epoch of a change, the live frontends and the workers that cannot follow. It reads the same report as GET /api/cluster/carrier, in process or over HTTP. No tool changes the carrier. A change moves every frontend and worker and needs a dry run and a check of the frontend list by an admin, so the skill tells the assistant to point the user to the dashboard or the CLI. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/application/application.go | 3 ++ .../endpoints/mcp/localai_assistant_test.go | 4 ++ pkg/mcp/localaitools/client.go | 6 +++ pkg/mcp/localaitools/coverage_test.go | 1 + pkg/mcp/localaitools/dto.go | 38 ++++++++++++++ pkg/mcp/localaitools/fakes_test.go | 5 ++ pkg/mcp/localaitools/httpapi/client.go | 50 +++++++++++++++++++ pkg/mcp/localaitools/httpapi/client_test.go | 37 ++++++++++++++ pkg/mcp/localaitools/httpapi/routes.go | 1 + pkg/mcp/localaitools/inproc/client.go | 41 +++++++++++++++ pkg/mcp/localaitools/inproc/client_test.go | 42 ++++++++++++++++ pkg/mcp/localaitools/prompts/20_tools.md | 1 + .../prompts/skills/system_status.md | 3 +- pkg/mcp/localaitools/server.go | 1 + pkg/mcp/localaitools/server_test.go | 2 + pkg/mcp/localaitools/tools.go | 5 ++ pkg/mcp/localaitools/tools_cluster.go | 25 ++++++++++ 17 files changed, 264 insertions(+), 1 deletion(-) create mode 100644 pkg/mcp/localaitools/tools_cluster.go diff --git a/core/application/application.go b/core/application/application.go index d22895af6ad7..b3ab5e15fb60 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -611,6 +611,9 @@ func (a *Application) start() error { if a.distributed != nil && a.distributed.Unloader != nil { assistantClient.LoadStopper = a.distributed.Unloader } + if a.distributed != nil && a.distributed.Switch != nil { + assistantClient.Carrier = a.distributed.Switch + } // PII filter — same nil-or-real wiring. assistantClient.PIIRedactor = a.piiRedactor assistantClient.PIIEvents = a.piiEvents diff --git a/core/http/endpoints/mcp/localai_assistant_test.go b/core/http/endpoints/mcp/localai_assistant_test.go index 3e539ce7d4bd..f6409218e89f 100644 --- a/core/http/endpoints/mcp/localai_assistant_test.go +++ b/core/http/endpoints/mcp/localai_assistant_test.go @@ -218,6 +218,10 @@ func (stubClient) ClearRouterCorpus(_ context.Context, routerModel string) (*loc return &localaitools.RouterCorpusClearResult{Router: routerModel}, nil } +func (stubClient) GetClusterCarrier(_ context.Context) (*localaitools.ClusterCarrierStatus, error) { + return &localaitools.ClusterCarrierStatus{}, nil +} + func (stubClient) ListFailoverChains(_ context.Context) ([]localaitools.FailoverChainInfo, error) { return []localaitools.FailoverChainInfo{}, nil } diff --git a/pkg/mcp/localaitools/client.go b/pkg/mcp/localaitools/client.go index ba5c2834a76f..6e4addcc89b7 100644 --- a/pkg/mcp/localaitools/client.go +++ b/pkg/mcp/localaitools/client.go @@ -134,6 +134,12 @@ type LocalAIClient interface { // index. ClearRouterCorpus(ctx context.Context, routerModel string) (*RouterCorpusClearResult, error) + // ---- Distributed carrier ---- + // GetClusterCarrier reports which transport the cluster uses (NATS or the + // database tunnel), the state of a change, the frontends and the workers + // that cannot follow a change. It never changes anything. + GetClusterCarrier(ctx context.Context) (*ClusterCarrierStatus, error) + // ---- Failover chains ---- // ListFailoverChains reports every configured failover chain, its // currently active target, and the health of each target. diff --git a/pkg/mcp/localaitools/coverage_test.go b/pkg/mcp/localaitools/coverage_test.go index 6ec5720249d3..87025e293d6c 100644 --- a/pkg/mcp/localaitools/coverage_test.go +++ b/pkg/mcp/localaitools/coverage_test.go @@ -47,6 +47,7 @@ var toolToHTTPRoute = map[string]string{ ToolListAliases: "GET /api/aliases", ToolListVoiceProfiles: "GET /api/voice-profiles", ToolListFailoverChains: "GET /api/failover", + ToolGetClusterCarrier: "GET /api/cluster/carrier", // Mutating tools. ToolInstallModel: "POST /models/apply", diff --git a/pkg/mcp/localaitools/dto.go b/pkg/mcp/localaitools/dto.go index 88d7159ffa8b..1a975b22254f 100644 --- a/pkg/mcp/localaitools/dto.go +++ b/pkg/mcp/localaitools/dto.go @@ -414,6 +414,44 @@ type VRAMEstimateRequest struct { KVQuantBits int `json:"kv_quant_bits,omitempty" jsonschema:"KV cache quantization bits (e.g. 4, 8, 16)."` } +// ClusterCarrierStatus is the LLM-facing view of the transport of a +// distributed cluster. Distributed is false on a frontend that is not +// distributed, and the other fields are then empty. +type ClusterCarrierStatus struct { + Distributed bool `json:"distributed"` + // Active is "nats" or "tunnel". + Active string `json:"active,omitempty"` + // State is "stable", "prepare" or "commit". A stable cluster with + // DrainRemainingSeconds above zero is draining the previous carrier. + State string `json:"state,omitempty"` + Target string `json:"target,omitempty"` + Epoch int64 `json:"epoch,omitempty"` + DrainRemainingSeconds float64 `json:"drain_remaining_seconds,omitempty"` + // Replicas are the live frontends. + Replicas []ClusterReplicaInfo `json:"replicas,omitempty"` + // Workers are the registered workers and the carriers they hold. + Workers []ClusterWorkerInfo `json:"workers,omitempty"` +} + +// ClusterReplicaInfo is one live frontend replica. +type ClusterReplicaInfo struct { + ID string `json:"id"` + Version string `json:"version,omitempty"` + ReadyEpoch int64 `json:"ready_epoch"` + ReadyReason string `json:"ready_reason,omitempty"` +} + +// ClusterWorkerInfo is one worker. CanFollow is about a change to the other +// carrier, and Reason says why it is false. +type ClusterWorkerInfo struct { + ID string `json:"id"` + Name string `json:"name,omitempty"` + Attached []string `json:"attached,omitempty"` + CanFollow bool `json:"can_follow"` + Reason string `json:"reason,omitempty"` + FollowError string `json:"follow_error,omitempty"` +} + // FailoverTargetInfo is the LLM-facing view of one failover chain target's // health. It mirrors failover.TargetStatus but drops ConsecutiveOK and // LastProbe — internal probing detail the LLM doesn't need to decide diff --git a/pkg/mcp/localaitools/fakes_test.go b/pkg/mcp/localaitools/fakes_test.go index d8a85491d721..2b5d5c0c52e8 100644 --- a/pkg/mcp/localaitools/fakes_test.go +++ b/pkg/mcp/localaitools/fakes_test.go @@ -402,6 +402,11 @@ func (f *fakeClient) ClearRouterCorpus(_ context.Context, routerModel string) (* return &RouterCorpusClearResult{Router: routerModel}, nil } +func (f *fakeClient) GetClusterCarrier(_ context.Context) (*ClusterCarrierStatus, error) { + f.record("GetClusterCarrier", nil) + return &ClusterCarrierStatus{Distributed: true, Active: "tunnel", State: "stable", Epoch: 1}, nil +} + func (f *fakeClient) ListFailoverChains(_ context.Context) ([]FailoverChainInfo, error) { f.record("ListFailoverChains", nil) if f.listFailoverChains != nil { diff --git a/pkg/mcp/localaitools/httpapi/client.go b/pkg/mcp/localaitools/httpapi/client.go index 726f121a30be..983e0986c1d6 100644 --- a/pkg/mcp/localaitools/httpapi/client.go +++ b/pkg/mcp/localaitools/httpapi/client.go @@ -830,6 +830,56 @@ func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*lo return &out, nil } +// ---- Cluster carrier ---- + +// GetClusterCarrier reads GET /api/cluster/carrier. A frontend that is not +// distributed answers 503, which is reported as distributed=false. +func (c *Client) GetClusterCarrier(ctx context.Context) (*localaitools.ClusterCarrierStatus, error) { + var raw struct { + Active string `json:"active"` + State string `json:"state"` + Target string `json:"target"` + Epoch int64 `json:"epoch"` + DrainRemaining int64 `json:"drain_remaining_ns"` + Replicas []struct { + ID string `json:"id"` + Version string `json:"version"` + ReadyEpoch int64 `json:"ready_epoch"` + ReadyReason string `json:"ready_reason"` + } `json:"replicas"` + Workers []struct { + ID string `json:"id"` + Name string `json:"name"` + Attached []string `json:"attached"` + CanFollow bool `json:"can_follow"` + Reason string `json:"reason"` + FollowError string `json:"follow_error"` + } `json:"workers"` + } + if err := c.do(ctx, http.MethodGet, routeClusterCarrier, nil, &raw); err != nil { + var he *HTTPError + if errors.As(err, &he) && he.StatusCode == http.StatusServiceUnavailable { + return &localaitools.ClusterCarrierStatus{Distributed: false}, nil + } + return nil, err + } + out := &localaitools.ClusterCarrierStatus{ + Distributed: true, + Active: raw.Active, + State: raw.State, + Target: raw.Target, + Epoch: raw.Epoch, + DrainRemainingSeconds: float64(raw.DrainRemaining) / 1e9, + } + for _, r := range raw.Replicas { + out.Replicas = append(out.Replicas, localaitools.ClusterReplicaInfo{ID: r.ID, Version: r.Version, ReadyEpoch: r.ReadyEpoch, ReadyReason: r.ReadyReason}) + } + for _, w := range raw.Workers { + out.Workers = append(out.Workers, localaitools.ClusterWorkerInfo{ID: w.ID, Name: w.Name, Attached: w.Attached, CanFollow: w.CanFollow, Reason: w.Reason, FollowError: w.FollowError}) + } + return out, nil +} + // ---- Failover chains ---- func (c *Client) ListFailoverChains(ctx context.Context) ([]localaitools.FailoverChainInfo, error) { diff --git a/pkg/mcp/localaitools/httpapi/client_test.go b/pkg/mcp/localaitools/httpapi/client_test.go index fd1cff632be7..cf1434880262 100644 --- a/pkg/mcp/localaitools/httpapi/client_test.go +++ b/pkg/mcp/localaitools/httpapi/client_test.go @@ -482,3 +482,40 @@ var _ = Describe("Bearer token", func() { Expect(sawAuth).To(Equal("Bearer secret-key")) }) }) + +var _ = Describe("Cluster carrier", func() { + It("issues GET /api/cluster/carrier and maps the report", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + Expect(r.Method).To(Equal(http.MethodGet)) + Expect(r.URL.Path).To(Equal("/api/cluster/carrier")) + _ = json.NewEncoder(w).Encode(map[string]any{ + "active": "nats", "state": "commit", "target": "tunnel", "epoch": 7, + "drain_remaining_ns": 30e9, + "replicas": []map[string]any{{"id": "r1", "version": "v1", "ready_epoch": 7}}, + "workers": []map[string]any{{"id": "w1", "name": "n", "attached": []string{"nats"}, + "can_follow": false, "reason": "old worker", "follow_error": "x"}}, + }) + })) + DeferCleanup(srv.Close) + + out, err := New(srv.URL, "").GetClusterCarrier(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(out.Distributed).To(BeTrue()) + Expect(out.Active).To(Equal("nats")) + Expect(out.Target).To(Equal("tunnel")) + Expect(out.DrainRemainingSeconds).To(BeNumerically("==", 30)) + Expect(out.Replicas).To(HaveLen(1)) + Expect(out.Workers[0].Reason).To(Equal("old worker")) + }) + + It("reports distributed=false when the server answers 503", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte(`{"error":"distributed mode is not enabled on this frontend"}`)) + })) + DeferCleanup(srv.Close) + out, err := New(srv.URL, "").GetClusterCarrier(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(out.Distributed).To(BeFalse()) + }) +}) diff --git a/pkg/mcp/localaitools/httpapi/routes.go b/pkg/mcp/localaitools/httpapi/routes.go index 229e6d0ff470..090f65eeacf9 100644 --- a/pkg/mcp/localaitools/httpapi/routes.go +++ b/pkg/mcp/localaitools/httpapi/routes.go @@ -35,6 +35,7 @@ const ( routeRouterDecisions = "/api/router/decisions" routeVoiceProfiles = "/api/voice-profiles" routeFailover = "/api/failover" + routeClusterCarrier = "/api/cluster/carrier" ) func routeJobStatus(jobID string) string { diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go index b1900656bb43..b274764a0bb8 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -21,6 +21,7 @@ import ( "github.com/mudler/LocalAI/core/gallery/importers" "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/failover" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/modeladmin" @@ -93,6 +94,10 @@ type Client struct { // configured. Failover *failover.Manager + // Carrier backs get_cluster_carrier. nil on a frontend that is not + // distributed, where the tool reports distributed=false. + Carrier CarrierStatusSource + modelAdmin *modeladmin.ConfigService } @@ -1114,6 +1119,42 @@ func (c *Client) ClearRouterCorpus(ctx context.Context, routerModel string) (*lo return &localaitools.RouterCorpusClearResult{Router: cfg.Name, Cleared: cleared}, nil } +// ---- Cluster carrier ---- + +// CarrierStatusSource is what get_cluster_carrier reads. *cluster.Switch is one. +type CarrierStatusSource interface { + Status(ctx context.Context) (cluster.Report, error) +} + +func (c *Client) GetClusterCarrier(ctx context.Context) (*localaitools.ClusterCarrierStatus, error) { + if c.Carrier == nil { + return &localaitools.ClusterCarrierStatus{Distributed: false}, nil + } + rep, err := c.Carrier.Status(ctx) + if err != nil { + return nil, err + } + out := &localaitools.ClusterCarrierStatus{ + Distributed: true, + Active: string(rep.Active), + State: string(rep.State), + Target: string(rep.Target), + Epoch: rep.Epoch, + DrainRemainingSeconds: rep.DrainRemaining.Seconds(), + } + for _, r := range rep.Replicas { + out.Replicas = append(out.Replicas, localaitools.ClusterReplicaInfo{ID: r.ID, Version: r.Version, ReadyEpoch: r.ReadyEpoch, ReadyReason: r.ReadyReason}) + } + for _, w := range rep.Workers { + attached := make([]string, 0, len(w.Attached)) + for _, a := range w.Attached { + attached = append(attached, string(a)) + } + out.Workers = append(out.Workers, localaitools.ClusterWorkerInfo{ID: w.ID, Name: w.Name, Attached: attached, CanFollow: w.CanFollow, Reason: w.Reason, FollowError: w.FollowError}) + } + return out, nil +} + // ---- Failover chains ---- func (c *Client) ListFailoverChains(_ context.Context) ([]localaitools.FailoverChainInfo, error) { diff --git a/pkg/mcp/localaitools/inproc/client_test.go b/pkg/mcp/localaitools/inproc/client_test.go index 0a9e6a5b761f..62c39d5b0312 100644 --- a/pkg/mcp/localaitools/inproc/client_test.go +++ b/pkg/mcp/localaitools/inproc/client_test.go @@ -13,6 +13,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/failover" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/nodes" @@ -303,3 +304,44 @@ func mustMarshal(v any) string { Expect(err).ToNot(HaveOccurred()) return string(b) } + +type fakeCarrierSource struct { + report cluster.Report + err error +} + +func (f fakeCarrierSource) Status(context.Context) (cluster.Report, error) { return f.report, f.err } + +var _ = Describe("GetClusterCarrier", func() { + It("reports distributed=false when no switch is wired", func() { + out, err := (&Client{}).GetClusterCarrier(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(out.Distributed).To(BeFalse()) + }) + + It("maps the report to the assistant view", func() { + c := &Client{Carrier: fakeCarrierSource{report: cluster.Report{ + Active: cluster.CarrierTunnel, State: cluster.StateStable, Epoch: 4, + DrainRemaining: 90 * time.Second, + Replicas: []cluster.ReplicaStatus{{ID: "r1", Version: "v1", ReadyEpoch: 4}}, + Workers: []cluster.WorkerStatus{{ID: "w1", Name: "gpu", Attached: []cluster.Carrier{cluster.CarrierTunnel}, + CanFollow: false, Reason: "no routable address", FollowError: "no address"}}, + }}} + out, err := c.GetClusterCarrier(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(out.Distributed).To(BeTrue()) + Expect(out.Active).To(Equal("tunnel")) + Expect(out.Epoch).To(Equal(int64(4))) + Expect(out.DrainRemainingSeconds).To(BeNumerically("==", 90)) + Expect(out.Replicas).To(ConsistOf(localaitools.ClusterReplicaInfo{ID: "r1", Version: "v1", ReadyEpoch: 4})) + Expect(out.Workers).To(HaveLen(1)) + Expect(out.Workers[0].Attached).To(Equal([]string{"tunnel"})) + Expect(out.Workers[0].CanFollow).To(BeFalse()) + Expect(out.Workers[0].FollowError).To(Equal("no address")) + }) + + It("passes an error from the switch on", func() { + _, err := (&Client{Carrier: fakeCarrierSource{err: errors.New("db down")}}).GetClusterCarrier(context.Background()) + Expect(err).To(MatchError(ContainSubstring("db down"))) + }) +}) diff --git a/pkg/mcp/localaitools/prompts/20_tools.md b/pkg/mcp/localaitools/prompts/20_tools.md index 591c49f51a40..82f26bec5875 100644 --- a/pkg/mcp/localaitools/prompts/20_tools.md +++ b/pkg/mcp/localaitools/prompts/20_tools.md @@ -25,6 +25,7 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the - `get_router_corpus_stats` — Inspect a KNN router corpus by count and label only; exemplar texts are never returned. - `list_aliases` — List configured model aliases and their targets. - `list_failover_chains` — List failover chains, their active target and target health. +- `get_cluster_carrier` — Read which transport a distributed cluster uses (NATS or the database tunnel), the state of a change, the live frontends and the workers that cannot follow. Read-only. No tool changes the carrier. ## Mutating (require user confirmation per safety rule 1) diff --git a/pkg/mcp/localaitools/prompts/skills/system_status.md b/pkg/mcp/localaitools/prompts/skills/system_status.md index 7f7146344a79..859f99ede96d 100644 --- a/pkg/mcp/localaitools/prompts/skills/system_status.md +++ b/pkg/mcp/localaitools/prompts/skills/system_status.md @@ -4,7 +4,7 @@ Use this when the user asks "what's installed?", "what's running?", "show status 1. Call `system_info` for version, paths, distributed flag, loaded models, installed backends. 2. Call `list_installed_models` (no capability filter) for the full installed-model inventory. -3. If `system_info.distributed` is true, also call `list_nodes` and `list_scheduling`, then report worker health and any per-model scheduling rules. +3. If `system_info.distributed` is true, also call `list_nodes`, `list_scheduling` and `get_cluster_carrier`, then report worker health, any per-model scheduling rules, and the transport (`active`, `state`, and any worker with `can_follow` false or a `follow_error`). You cannot change the transport: no tool does it. If the user asks for a change, tell them to use the Cluster transport panel in the Nodes page of the dashboard or `local-ai cluster carrier switch`, because the change needs a dry run and a check of the frontend list by an admin. 4. Present a concise summary: - **Version & mode** (`distributed: true|false`) - **Installed models** (count + list, each with name and capabilities) @@ -12,4 +12,5 @@ Use this when the user asks "what's installed?", "what's running?", "show status - **Loaded right now** (from `loaded_models`) - **Workers** (only when distributed) - **Scheduling rules** (only when distributed) + - **Transport** (only when distributed) 5. Do not call mutating tools in this skill. diff --git a/pkg/mcp/localaitools/server.go b/pkg/mcp/localaitools/server.go index f47332db5203..3af8b5357b5f 100644 --- a/pkg/mcp/localaitools/server.go +++ b/pkg/mcp/localaitools/server.go @@ -55,6 +55,7 @@ func NewServer(client LocalAIClient, opts Options) *mcp.Server { registerPIITools(srv, client, opts) registerMiddlewareTools(srv, client, opts) registerFailoverTools(srv, client, opts) + registerClusterTools(srv, client) return srv } diff --git a/pkg/mcp/localaitools/server_test.go b/pkg/mcp/localaitools/server_test.go index cc33024ec00e..23da55a82183 100644 --- a/pkg/mcp/localaitools/server_test.go +++ b/pkg/mcp/localaitools/server_test.go @@ -93,6 +93,7 @@ var expectedReadOnlyCatalog = sortedStrings( ToolSystemInfo, ToolVRAMEstimate, ToolListFailoverChains, + ToolGetClusterCarrier, ) // expectedFullCatalog derives from the read-only catalog plus the canonical @@ -158,6 +159,7 @@ var _ = Describe("Tool dispatch", func() { {ToolCreateVoiceProfile, CreateVoiceProfileRequest{Name: "Narrator", Transcript: "Reference words", AudioBase64: "UklGRg==", ConsentConfirmed: true}, "CreateVoiceProfile"}, {ToolDeleteVoiceProfile, DeleteVoiceProfileRequest{ID: "00000000-0000-0000-0000-000000000001"}, "DeleteVoiceProfile"}, {ToolListFailoverChains, map[string]any{}, "ListFailoverChains"}, + {ToolGetClusterCarrier, map[string]any{}, "GetClusterCarrier"}, {ToolPinFailoverTarget, map[string]any{"chain": "c", "target": "b"}, "PinFailoverTarget"}, {ToolUnpinFailoverTarget, map[string]any{"chain": "c"}, "UnpinFailoverTarget"}, } diff --git a/pkg/mcp/localaitools/tools.go b/pkg/mcp/localaitools/tools.go index ee03d196e631..5e348a2d59f1 100644 --- a/pkg/mcp/localaitools/tools.go +++ b/pkg/mcp/localaitools/tools.go @@ -63,6 +63,11 @@ const ( // ToolListFailoverChains is read-only but lives here so the failover // tools stay grouped; the catalog tests assert its read-only placement. ToolListFailoverChains = "list_failover_chains" + + // ToolGetClusterCarrier is read-only. There is no tool that changes the + // carrier: a change moves every frontend and worker of the cluster, so an + // admin makes it in the dashboard or with the `local-ai cluster` command. + ToolGetClusterCarrier = "get_cluster_carrier" ) // DefaultServerName is the MCP Implementation.Name surfaced when diff --git a/pkg/mcp/localaitools/tools_cluster.go b/pkg/mcp/localaitools/tools_cluster.go new file mode 100644 index 000000000000..cd48572d4aec --- /dev/null +++ b/pkg/mcp/localaitools/tools_cluster.go @@ -0,0 +1,25 @@ +package localaitools + +import ( + "context" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// registerClusterTools wires the read-only status of the cluster carrier. +// A change of carrier has no tool on purpose. It moves every frontend and +// worker, and it needs a dry run and a list of the frontends that only an +// admin can confirm, so the assistant points to the dashboard or to the +// `local-ai cluster carrier` command. +func registerClusterTools(s *mcp.Server, client LocalAIClient) { + mcp.AddTool(s, &mcp.Tool{ + Name: ToolGetClusterCarrier, + Description: "Read which transport a distributed cluster uses (NATS or the database tunnel), the state and epoch of a change, the live frontends and the workers that cannot follow a change. Read-only. Returns distributed=false when this server is not distributed.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + status, err := client.GetClusterCarrier(ctx) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(status), nil, nil + }) +} From 92c9e51c8d55ebe47a8c1a6697888d0174099827 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:52:40 +0000 Subject: [PATCH 148/162] docs(api): document the cluster carrier and settings endpoints in swagger Add swagger blocks to the four admin routes under /api/cluster/ and regenerate the spec. Export the request and settings types so the spec can name them, and mark the duration fields as integers (nanoseconds) so the generator can read them. The JSON that the routes send does not change. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../http/endpoints/localai/cluster_carrier.go | 85 ++- core/services/cluster/settings.go | 6 +- core/services/cluster/switch.go | 4 +- swagger/docs.go | 548 ++++++++++++++++++ swagger/swagger.json | 548 ++++++++++++++++++ swagger/swagger.yaml | 412 +++++++++++++ 6 files changed, 1582 insertions(+), 21 deletions(-) diff --git a/core/http/endpoints/localai/cluster_carrier.go b/core/http/endpoints/localai/cluster_carrier.go index 1a3f2692aa2b..ac8b4cbb63b7 100644 --- a/core/http/endpoints/localai/cluster_carrier.go +++ b/core/http/endpoints/localai/cluster_carrier.go @@ -69,27 +69,40 @@ func carrierError(c echo.Context, status int, msg string, extra map[string]any) // settingsView is the settings as an admin reads them. A user name, a password or a // query string in the address of a server is dropped. -func settingsView(all map[string]string) map[string]any { - view := map[string]any{ - "nats_url": "", - "nats_worker_url": "", - "prepare_timeout": all[cluster.SettingPrepareTimeout], - "transition_window": all[cluster.SettingTransitionWindow], - "max_drain": all[cluster.SettingMaxDrain], +func settingsView(all map[string]string) ClusterSettingsView { + view := ClusterSettingsView{ + PrepareTimeout: all[cluster.SettingPrepareTimeout], + TransitionWindow: all[cluster.SettingTransitionWindow], + MaxDrain: all[cluster.SettingMaxDrain], } if v := all[cluster.SettingNATSURL]; v != "" { - view["nats_url"] = cluster.PublicNATSURL(v) + view.NATSURL = cluster.PublicNATSURL(v) } if v := all[cluster.SettingNATSWorkerURL]; v != "" { - view["nats_worker_url"] = cluster.PublicNATSURL(v) + view.NATSWorkerURL = cluster.PublicNATSURL(v) } return view } +// ClusterSettingsView is the settings of the cluster as an admin reads them. A +// value that is not stored is empty and the default applies. +type ClusterSettingsView struct { + // NATSURL is the NATS address that the frontends use. + NATSURL string `json:"nats_url"` + // NATSWorkerURL is the NATS address that workers are told to use. Empty means NATSURL. + NATSWorkerURL string `json:"nats_worker_url"` + // PrepareTimeout is how long a change waits for every replica to be ready (a Go duration). + PrepareTimeout string `json:"prepare_timeout"` + // TransitionWindow is how long a replica may take to confirm a commit (a Go duration). + TransitionWindow string `json:"transition_window"` + // MaxDrain is how long the previous carrier stays attached after a commit (a Go duration). + MaxDrain string `json:"max_drain"` +} + // carrierView is the report with the settings next to it. -type carrierView struct { +type CarrierStatusResponse struct { cluster.Report - Settings map[string]any `json:"settings"` + Settings ClusterSettingsView `json:"settings"` } // GetCarrierEndpoint reports the state of the carrier of the cluster: the active @@ -97,6 +110,13 @@ type carrierView struct { // the workers that could not follow a change to the other carrier and why, the // work in flight, and the list of live replicas for the admin to confirm before a // change. +// +// @Summary Report the transport of a distributed cluster +// @Description Returns the active carrier (nats or tunnel), the epoch, the state of a change (stable, prepare or commit), the live replicas with their readiness, the workers and whether each could follow a change to the other carrier, the work in flight, and the cluster settings. Admin only. Answers 503 when distributed mode is off. +// @Tags Nodes +// @Success 200 {object} localai.CarrierStatusResponse +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/carrier [get] func GetCarrierEndpoint(sw CarrierSwitchAdmin, settings ClusterSettings) echo.HandlerFunc { return func(c echo.Context) error { ctx := c.Request().Context() @@ -118,11 +138,12 @@ func GetCarrierEndpoint(sw CarrierSwitchAdmin, settings ClusterSettings) echo.Ha if err != nil { return carrierError(c, http.StatusInternalServerError, err.Error(), nil) } - return c.JSON(http.StatusOK, carrierView{Report: report, Settings: settingsView(all)}) + return c.JSON(http.StatusOK, CarrierStatusResponse{Report: report, Settings: settingsView(all)}) } } -type switchRequest struct { +// CarrierSwitchRequest is the body of POST /api/cluster/carrier. +type CarrierSwitchRequest struct { Target string `json:"target"` DryRun bool `json:"dry_run"` Force bool `json:"force"` @@ -138,9 +159,21 @@ type switchRequest struct { // when nothing blocks. The change itself is carried out by the replicas and led by // one of them, so the answer is 202: it says what was started and the epoch to // watch. +// +// @Summary Dry-run, start or abort a change of carrier +// @Description Takes a target (nats or tunnel) with dry_run and force, or abort. A dry run returns the preflight report and changes nothing. A request starts the change and returns 202 with the epoch to watch. A request that a blocker stops returns 422 with the blockers; force accepts the blockers that are forceable. A change already under way returns 409, and only abort is accepted then. Admin only. +// @Tags Nodes +// @Param request body localai.CarrierSwitchRequest true "target, dry_run, force or abort" +// @Success 200 {object} cluster.Report "Dry run report, or the row after an abort" +// @Success 202 {object} map[string]interface{} "Change started: state, active, target, epoch, force, warnings, row" +// @Failure 400 {object} map[string]string "Bad body or target" +// @Failure 409 {object} map[string]string "A change is under way, or there is nothing to abort" +// @Failure 422 {object} map[string]interface{} "Blocked: error, blockers, report" +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/carrier [post] func SwitchCarrierEndpoint(sw CarrierSwitchAdmin, prober ReplicaProber) echo.HandlerFunc { return func(c echo.Context) error { - var req switchRequest + var req CarrierSwitchRequest if err := c.Bind(&req); err != nil { return carrierError(c, http.StatusBadRequest, "invalid request body", nil) } @@ -197,6 +230,13 @@ func SwitchCarrierEndpoint(sw CarrierSwitchAdmin, prober ReplicaProber) echo.Han } // GetClusterSettingsEndpoint reads the settings of the cluster. +// +// @Summary Read the settings of the cluster +// @Description Returns the NATS address of the frontends and of the workers and the waits of a change. A user name, password or query string in an address is dropped. Admin only. +// @Tags Nodes +// @Success 200 {object} localai.ClusterSettingsView +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/settings [get] func GetClusterSettingsEndpoint(settings ClusterSettings) echo.HandlerFunc { return func(c echo.Context) error { all, err := settings.All(c.Request().Context()) @@ -207,7 +247,9 @@ func GetClusterSettingsEndpoint(settings ClusterSettings) echo.HandlerFunc { } } -type settingsRequest struct { +// ClusterSettingsRequest is the body of PUT /api/cluster/settings. A field that +// is absent is left as it is, and an empty string clears the setting. +type ClusterSettingsRequest struct { NATSURL *string `json:"nats_url"` NATSWorkerURL *string `json:"nats_worker_url"` PrepareTimeout *string `json:"prepare_timeout"` @@ -237,9 +279,20 @@ type CarrierStatus interface { // // No address may carry a user name, a password or a query string: the settings // are shared, and the credentials of NATS stay on each replica. +// +// @Summary Store settings of the cluster +// @Description Stores the NATS address, the address workers use, or the waits of a change. Saving a NATS address checks that the serving replica reaches it and makes NATS available; it does not switch the cluster. Admin only. +// @Tags Nodes +// @Param request body localai.ClusterSettingsRequest true "settings to store" +// @Success 200 {object} map[string]interface{} "saved, nats.reachable when a NATS address was checked, settings" +// @Failure 400 {object} map[string]string "Unknown setting, bad value, or a credential in an address" +// @Failure 409 {object} map[string]string "A change of carrier is under way" +// @Failure 422 {object} map[string]string "This replica cannot reach the NATS server" +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/settings [put] func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, prober ReplicaProber, state CarrierStatus) echo.HandlerFunc { return func(c echo.Context) error { - var req settingsRequest + var req ClusterSettingsRequest if err := c.Bind(&req); err != nil { return carrierError(c, http.StatusBadRequest, "invalid request body", nil) } diff --git a/core/services/cluster/settings.go b/core/services/cluster/settings.go index cd0cb2cd0bad..1e30bb347d5c 100644 --- a/core/services/cluster/settings.go +++ b/core/services/cluster/settings.go @@ -68,13 +68,13 @@ func CheckSetting(key, value string) error { return checkSetting(key, value) } // Timings are the waits of a change. type Timings struct { // PrepareTimeout is how long prepare waits for every live replica. - PrepareTimeout time.Duration + PrepareTimeout time.Duration `swaggertype:"integer"` // TransitionWindow is how long commit waits for every live replica to // confirm. - TransitionWindow time.Duration + TransitionWindow time.Duration `swaggertype:"integer"` // MaxDrain is how long the previous carrier stays attached after the commit. // Work that is still running on it then is failed by the reaper. - MaxDrain time.Duration + MaxDrain time.Duration `swaggertype:"integer"` } // DefaultTimings are used for a timing that nothing sets. diff --git a/core/services/cluster/switch.go b/core/services/cluster/switch.go index ca7f82224f01..0ec5c213a748 100644 --- a/core/services/cluster/switch.go +++ b/core/services/cluster/switch.go @@ -122,7 +122,7 @@ type ReplicaStatus struct { // Availability maps each carrier to the reason the replica cannot build it, // empty when it can. A carrier it never looked at is absent. Availability map[Carrier]string `json:"availability,omitempty"` - AvailabilityAge time.Duration `json:"availability_age_ns,omitempty"` + AvailabilityAge time.Duration `json:"availability_age_ns,omitempty" swaggertype:"integer"` } // WorkerStatus is one worker, and whether it can follow the target. @@ -155,7 +155,7 @@ type Report struct { Timings Timings `json:"timings"` // DrainRemaining is how long the previous carrier stays attached, when one // is draining. - DrainRemaining time.Duration `json:"drain_remaining_ns,omitempty"` + DrainRemaining time.Duration `json:"drain_remaining_ns,omitempty" swaggertype:"integer"` } // Request asks for a change of carrier. diff --git a/swagger/docs.go b/swagger/docs.go index 063447074be4..427c1681ef9c 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -881,6 +881,188 @@ const docTemplate = `{ } } }, + "/api/cluster/carrier": { + "get": { + "description": "Returns the active carrier (nats or tunnel), the epoch, the state of a change (stable, prepare or commit), the live replicas with their readiness, the workers and whether each could follow a change to the other carrier, the work in flight, and the cluster settings. Admin only. Answers 503 when distributed mode is off.", + "tags": [ + "Nodes" + ], + "summary": "Report the transport of a distributed cluster", + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/localai.CarrierStatusResponse" + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + }, + "post": { + "description": "Takes a target (nats or tunnel) with dry_run and force, or abort. A dry run returns the preflight report and changes nothing. A request starts the change and returns 202 with the epoch to watch. A request that a blocker stops returns 422 with the blockers; force accepts the blockers that are forceable. A change already under way returns 409, and only abort is accepted then. Admin only.", + "tags": [ + "Nodes" + ], + "summary": "Dry-run, start or abort a change of carrier", + "parameters": [ + { + "description": "target, dry_run, force or abort", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/localai.CarrierSwitchRequest" + } + } + ], + "responses": { + "200": { + "description": "Dry run report, or the row after an abort", + "schema": { + "$ref": "#/definitions/cluster.Report" + } + }, + "202": { + "description": "Change started: state, active, target, epoch, force, warnings, row", + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "400": { + "description": "Bad body or target", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "409": { + "description": "A change is under way, or there is nothing to abort", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "422": { + "description": "Blocked: error, blockers, report", + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, + "/api/cluster/settings": { + "get": { + "description": "Returns the NATS address of the frontends and of the workers and the waits of a change. A user name, password or query string in an address is dropped. Admin only.", + "tags": [ + "Nodes" + ], + "summary": "Read the settings of the cluster", + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/localai.ClusterSettingsView" + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + }, + "put": { + "description": "Stores the NATS address, the address workers use, or the waits of a change. Saving a NATS address checks that the serving replica reaches it and makes NATS available; it does not switch the cluster. Admin only.", + "tags": [ + "Nodes" + ], + "summary": "Store settings of the cluster", + "parameters": [ + { + "description": "settings to store", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/localai.ClusterSettingsRequest" + } + } + ], + "responses": { + "200": { + "description": "saved, nats.reachable when a NATS address was checked, settings", + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "400": { + "description": "Unknown setting, bad value, or a credential in an address", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "409": { + "description": "A change of carrier is under way", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "422": { + "description": "This replica cannot reach the NATS server", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, "/api/failover": { "get": { "produces": [ @@ -4446,6 +4628,250 @@ const docTemplate = `{ } }, "definitions": { + "cluster.Blocker": { + "type": "object", + "properties": { + "forceable": { + "description": "Forceable is true when a forced request goes ahead anyway.", + "type": "boolean" + }, + "id": { + "type": "string" + }, + "kind": { + "type": "string" + }, + "reason": { + "type": "string" + } + } + }, + "cluster.Carrier": { + "type": "string", + "enum": [ + "nats", + "tunnel" + ], + "x-enum-varnames": [ + "CarrierNATS", + "CarrierTunnel" + ] + }, + "cluster.CarrierRow": { + "type": "object", + "properties": { + "active": { + "$ref": "#/definitions/cluster.Carrier" + }, + "changed_at": { + "description": "ChangedAt is stamped by the database on every change, and the age of a\nchange is read on that same clock (CarrierStore.DBNow), so that the\ntimeouts of a change need no clock that the replicas agree on.", + "type": "string" + }, + "changed_by": { + "type": "string" + }, + "draining": { + "$ref": "#/definitions/cluster.Carrier" + }, + "draining_until": { + "type": "string" + }, + "epoch": { + "type": "integer" + }, + "force": { + "description": "Force records that the admin who started the change accepted that a\nreplica which is not ready, or a worker that cannot follow, is left\nbehind. It belongs to one change and is cleared by the next one.", + "type": "boolean" + }, + "nats_url": { + "description": "NATSURL is the address of the NATS server that this change, or the last\nchange to NATS, builds from. The admin starts a change and the replicas\nbuild in it, and an edit of the setting in between must not give two\nreplicas two servers, so every replica builds from this copy while a change\nis under way. It never carries a credential. Empty when the cluster has no\nNATS address.", + "type": "string" + }, + "note": { + "description": "Note says in a sentence how the last change ended, or why it did. It is\nfor operators.", + "type": "string" + }, + "prev_epoch": { + "type": "integer" + }, + "state": { + "$ref": "#/definitions/cluster.State" + }, + "target": { + "$ref": "#/definitions/cluster.Carrier" + } + } + }, + "cluster.InFlight": { + "type": "object", + "properties": { + "claimed_claims": { + "type": "integer" + }, + "jobs": { + "description": "Jobs is the number of jobs that run.", + "type": "integer" + }, + "loads": { + "description": "Loads is the number of models that are loading.", + "type": "integer" + }, + "pending_claims": { + "description": "PendingClaims and ClaimedClaims count the rows of the claim queue.", + "type": "integer" + } + } + }, + "cluster.ReplicaStatus": { + "type": "object", + "properties": { + "availability": { + "description": "Availability maps each carrier to the reason the replica cannot build it,\nempty when it can. A carrier it never looked at is absent.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "availability_age_ns": { + "type": "integer" + }, + "id": { + "type": "string" + }, + "ready_epoch": { + "type": "integer" + }, + "ready_reason": { + "type": "string" + }, + "version": { + "type": "string" + } + } + }, + "cluster.Report": { + "type": "object", + "properties": { + "active": { + "$ref": "#/definitions/cluster.Carrier" + }, + "blockers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Blocker" + } + }, + "drain_remaining_ns": { + "description": "DrainRemaining is how long the previous carrier stays attached, when one\nis draining.", + "type": "integer" + }, + "epoch": { + "type": "integer" + }, + "in_flight": { + "$ref": "#/definitions/cluster.InFlight" + }, + "ok": { + "type": "boolean" + }, + "replicas": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.ReplicaStatus" + } + }, + "row": { + "$ref": "#/definitions/cluster.CarrierRow" + }, + "state": { + "$ref": "#/definitions/cluster.State" + }, + "target": { + "$ref": "#/definitions/cluster.Carrier" + }, + "timings": { + "$ref": "#/definitions/cluster.Timings" + }, + "warnings": { + "type": "array", + "items": { + "type": "string" + } + }, + "workers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.WorkerStatus" + } + } + } + }, + "cluster.State": { + "type": "string", + "enum": [ + "stable", + "prepare", + "commit" + ], + "x-enum-varnames": [ + "StateStable", + "StatePrepare", + "StateCommit" + ] + }, + "cluster.Timings": { + "type": "object", + "properties": { + "maxDrain": { + "description": "MaxDrain is how long the previous carrier stays attached after the commit.\nWork that is still running on it then is failed by the reaper.", + "type": "integer" + }, + "prepareTimeout": { + "description": "PrepareTimeout is how long prepare waits for every live replica.", + "type": "integer" + }, + "transitionWindow": { + "description": "TransitionWindow is how long commit waits for every live replica to\nconfirm.", + "type": "integer" + } + } + }, + "cluster.WorkerStatus": { + "type": "object", + "properties": { + "attached": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Carrier" + } + }, + "can_follow": { + "description": "CanFollow and Reason are set when the report has a target.", + "type": "boolean" + }, + "follow": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Carrier" + } + }, + "follow_error": { + "type": "string" + }, + "id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "reason": { + "type": "string" + }, + "type": { + "type": "string" + } + } + }, "config.Gallery": { "type": "object", "properties": { @@ -5089,6 +5515,128 @@ const docTemplate = `{ } } }, + "localai.CarrierStatusResponse": { + "type": "object", + "properties": { + "active": { + "$ref": "#/definitions/cluster.Carrier" + }, + "blockers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Blocker" + } + }, + "drain_remaining_ns": { + "description": "DrainRemaining is how long the previous carrier stays attached, when one\nis draining.", + "type": "integer" + }, + "epoch": { + "type": "integer" + }, + "in_flight": { + "$ref": "#/definitions/cluster.InFlight" + }, + "ok": { + "type": "boolean" + }, + "replicas": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.ReplicaStatus" + } + }, + "row": { + "$ref": "#/definitions/cluster.CarrierRow" + }, + "settings": { + "$ref": "#/definitions/localai.ClusterSettingsView" + }, + "state": { + "$ref": "#/definitions/cluster.State" + }, + "target": { + "$ref": "#/definitions/cluster.Carrier" + }, + "timings": { + "$ref": "#/definitions/cluster.Timings" + }, + "warnings": { + "type": "array", + "items": { + "type": "string" + } + }, + "workers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.WorkerStatus" + } + } + } + }, + "localai.CarrierSwitchRequest": { + "type": "object", + "properties": { + "abort": { + "type": "boolean" + }, + "dry_run": { + "type": "boolean" + }, + "force": { + "type": "boolean" + }, + "target": { + "type": "string" + } + } + }, + "localai.ClusterSettingsRequest": { + "type": "object", + "properties": { + "max_drain": { + "type": "string" + }, + "nats_url": { + "type": "string" + }, + "nats_worker_url": { + "type": "string" + }, + "prepare_timeout": { + "type": "string" + }, + "transition_window": { + "type": "string" + } + } + }, + "localai.ClusterSettingsView": { + "type": "object", + "properties": { + "max_drain": { + "description": "MaxDrain is how long the previous carrier stays attached after a commit (a Go duration).", + "type": "string" + }, + "nats_url": { + "description": "NATSURL is the NATS address that the frontends use.", + "type": "string" + }, + "nats_worker_url": { + "description": "NATSWorkerURL is the NATS address that workers are told to use. Empty means NATSURL.", + "type": "string" + }, + "prepare_timeout": { + "description": "PrepareTimeout is how long a change waits for every replica to be ready (a Go duration).", + "type": "string" + }, + "transition_window": { + "description": "TransitionWindow is how long a replica may take to confirm a commit (a Go duration).", + "type": "string" + } + } + }, "localai.FailoverChainsResponse": { "type": "object", "properties": { diff --git a/swagger/swagger.json b/swagger/swagger.json index 7ccfd6984462..351c6689da90 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -878,6 +878,188 @@ } } }, + "/api/cluster/carrier": { + "get": { + "description": "Returns the active carrier (nats or tunnel), the epoch, the state of a change (stable, prepare or commit), the live replicas with their readiness, the workers and whether each could follow a change to the other carrier, the work in flight, and the cluster settings. Admin only. Answers 503 when distributed mode is off.", + "tags": [ + "Nodes" + ], + "summary": "Report the transport of a distributed cluster", + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/localai.CarrierStatusResponse" + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + }, + "post": { + "description": "Takes a target (nats or tunnel) with dry_run and force, or abort. A dry run returns the preflight report and changes nothing. A request starts the change and returns 202 with the epoch to watch. A request that a blocker stops returns 422 with the blockers; force accepts the blockers that are forceable. A change already under way returns 409, and only abort is accepted then. Admin only.", + "tags": [ + "Nodes" + ], + "summary": "Dry-run, start or abort a change of carrier", + "parameters": [ + { + "description": "target, dry_run, force or abort", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/localai.CarrierSwitchRequest" + } + } + ], + "responses": { + "200": { + "description": "Dry run report, or the row after an abort", + "schema": { + "$ref": "#/definitions/cluster.Report" + } + }, + "202": { + "description": "Change started: state, active, target, epoch, force, warnings, row", + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "400": { + "description": "Bad body or target", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "409": { + "description": "A change is under way, or there is nothing to abort", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "422": { + "description": "Blocked: error, blockers, report", + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, + "/api/cluster/settings": { + "get": { + "description": "Returns the NATS address of the frontends and of the workers and the waits of a change. A user name, password or query string in an address is dropped. Admin only.", + "tags": [ + "Nodes" + ], + "summary": "Read the settings of the cluster", + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/localai.ClusterSettingsView" + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + }, + "put": { + "description": "Stores the NATS address, the address workers use, or the waits of a change. Saving a NATS address checks that the serving replica reaches it and makes NATS available; it does not switch the cluster. Admin only.", + "tags": [ + "Nodes" + ], + "summary": "Store settings of the cluster", + "parameters": [ + { + "description": "settings to store", + "name": "request", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/localai.ClusterSettingsRequest" + } + } + ], + "responses": { + "200": { + "description": "saved, nats.reachable when a NATS address was checked, settings", + "schema": { + "type": "object", + "additionalProperties": true + } + }, + "400": { + "description": "Unknown setting, bad value, or a credential in an address", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "409": { + "description": "A change of carrier is under way", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "422": { + "description": "This replica cannot reach the NATS server", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + }, + "503": { + "description": "Distributed mode is not enabled", + "schema": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + } + } + }, "/api/failover": { "get": { "produces": [ @@ -4443,6 +4625,250 @@ } }, "definitions": { + "cluster.Blocker": { + "type": "object", + "properties": { + "forceable": { + "description": "Forceable is true when a forced request goes ahead anyway.", + "type": "boolean" + }, + "id": { + "type": "string" + }, + "kind": { + "type": "string" + }, + "reason": { + "type": "string" + } + } + }, + "cluster.Carrier": { + "type": "string", + "enum": [ + "nats", + "tunnel" + ], + "x-enum-varnames": [ + "CarrierNATS", + "CarrierTunnel" + ] + }, + "cluster.CarrierRow": { + "type": "object", + "properties": { + "active": { + "$ref": "#/definitions/cluster.Carrier" + }, + "changed_at": { + "description": "ChangedAt is stamped by the database on every change, and the age of a\nchange is read on that same clock (CarrierStore.DBNow), so that the\ntimeouts of a change need no clock that the replicas agree on.", + "type": "string" + }, + "changed_by": { + "type": "string" + }, + "draining": { + "$ref": "#/definitions/cluster.Carrier" + }, + "draining_until": { + "type": "string" + }, + "epoch": { + "type": "integer" + }, + "force": { + "description": "Force records that the admin who started the change accepted that a\nreplica which is not ready, or a worker that cannot follow, is left\nbehind. It belongs to one change and is cleared by the next one.", + "type": "boolean" + }, + "nats_url": { + "description": "NATSURL is the address of the NATS server that this change, or the last\nchange to NATS, builds from. The admin starts a change and the replicas\nbuild in it, and an edit of the setting in between must not give two\nreplicas two servers, so every replica builds from this copy while a change\nis under way. It never carries a credential. Empty when the cluster has no\nNATS address.", + "type": "string" + }, + "note": { + "description": "Note says in a sentence how the last change ended, or why it did. It is\nfor operators.", + "type": "string" + }, + "prev_epoch": { + "type": "integer" + }, + "state": { + "$ref": "#/definitions/cluster.State" + }, + "target": { + "$ref": "#/definitions/cluster.Carrier" + } + } + }, + "cluster.InFlight": { + "type": "object", + "properties": { + "claimed_claims": { + "type": "integer" + }, + "jobs": { + "description": "Jobs is the number of jobs that run.", + "type": "integer" + }, + "loads": { + "description": "Loads is the number of models that are loading.", + "type": "integer" + }, + "pending_claims": { + "description": "PendingClaims and ClaimedClaims count the rows of the claim queue.", + "type": "integer" + } + } + }, + "cluster.ReplicaStatus": { + "type": "object", + "properties": { + "availability": { + "description": "Availability maps each carrier to the reason the replica cannot build it,\nempty when it can. A carrier it never looked at is absent.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "availability_age_ns": { + "type": "integer" + }, + "id": { + "type": "string" + }, + "ready_epoch": { + "type": "integer" + }, + "ready_reason": { + "type": "string" + }, + "version": { + "type": "string" + } + } + }, + "cluster.Report": { + "type": "object", + "properties": { + "active": { + "$ref": "#/definitions/cluster.Carrier" + }, + "blockers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Blocker" + } + }, + "drain_remaining_ns": { + "description": "DrainRemaining is how long the previous carrier stays attached, when one\nis draining.", + "type": "integer" + }, + "epoch": { + "type": "integer" + }, + "in_flight": { + "$ref": "#/definitions/cluster.InFlight" + }, + "ok": { + "type": "boolean" + }, + "replicas": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.ReplicaStatus" + } + }, + "row": { + "$ref": "#/definitions/cluster.CarrierRow" + }, + "state": { + "$ref": "#/definitions/cluster.State" + }, + "target": { + "$ref": "#/definitions/cluster.Carrier" + }, + "timings": { + "$ref": "#/definitions/cluster.Timings" + }, + "warnings": { + "type": "array", + "items": { + "type": "string" + } + }, + "workers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.WorkerStatus" + } + } + } + }, + "cluster.State": { + "type": "string", + "enum": [ + "stable", + "prepare", + "commit" + ], + "x-enum-varnames": [ + "StateStable", + "StatePrepare", + "StateCommit" + ] + }, + "cluster.Timings": { + "type": "object", + "properties": { + "maxDrain": { + "description": "MaxDrain is how long the previous carrier stays attached after the commit.\nWork that is still running on it then is failed by the reaper.", + "type": "integer" + }, + "prepareTimeout": { + "description": "PrepareTimeout is how long prepare waits for every live replica.", + "type": "integer" + }, + "transitionWindow": { + "description": "TransitionWindow is how long commit waits for every live replica to\nconfirm.", + "type": "integer" + } + } + }, + "cluster.WorkerStatus": { + "type": "object", + "properties": { + "attached": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Carrier" + } + }, + "can_follow": { + "description": "CanFollow and Reason are set when the report has a target.", + "type": "boolean" + }, + "follow": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Carrier" + } + }, + "follow_error": { + "type": "string" + }, + "id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "reason": { + "type": "string" + }, + "type": { + "type": "string" + } + } + }, "config.Gallery": { "type": "object", "properties": { @@ -5086,6 +5512,128 @@ } } }, + "localai.CarrierStatusResponse": { + "type": "object", + "properties": { + "active": { + "$ref": "#/definitions/cluster.Carrier" + }, + "blockers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.Blocker" + } + }, + "drain_remaining_ns": { + "description": "DrainRemaining is how long the previous carrier stays attached, when one\nis draining.", + "type": "integer" + }, + "epoch": { + "type": "integer" + }, + "in_flight": { + "$ref": "#/definitions/cluster.InFlight" + }, + "ok": { + "type": "boolean" + }, + "replicas": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.ReplicaStatus" + } + }, + "row": { + "$ref": "#/definitions/cluster.CarrierRow" + }, + "settings": { + "$ref": "#/definitions/localai.ClusterSettingsView" + }, + "state": { + "$ref": "#/definitions/cluster.State" + }, + "target": { + "$ref": "#/definitions/cluster.Carrier" + }, + "timings": { + "$ref": "#/definitions/cluster.Timings" + }, + "warnings": { + "type": "array", + "items": { + "type": "string" + } + }, + "workers": { + "type": "array", + "items": { + "$ref": "#/definitions/cluster.WorkerStatus" + } + } + } + }, + "localai.CarrierSwitchRequest": { + "type": "object", + "properties": { + "abort": { + "type": "boolean" + }, + "dry_run": { + "type": "boolean" + }, + "force": { + "type": "boolean" + }, + "target": { + "type": "string" + } + } + }, + "localai.ClusterSettingsRequest": { + "type": "object", + "properties": { + "max_drain": { + "type": "string" + }, + "nats_url": { + "type": "string" + }, + "nats_worker_url": { + "type": "string" + }, + "prepare_timeout": { + "type": "string" + }, + "transition_window": { + "type": "string" + } + } + }, + "localai.ClusterSettingsView": { + "type": "object", + "properties": { + "max_drain": { + "description": "MaxDrain is how long the previous carrier stays attached after a commit (a Go duration).", + "type": "string" + }, + "nats_url": { + "description": "NATSURL is the NATS address that the frontends use.", + "type": "string" + }, + "nats_worker_url": { + "description": "NATSWorkerURL is the NATS address that workers are told to use. Empty means NATSURL.", + "type": "string" + }, + "prepare_timeout": { + "description": "PrepareTimeout is how long a change waits for every replica to be ready (a Go duration).", + "type": "string" + }, + "transition_window": { + "description": "TransitionWindow is how long a replica may take to confirm a commit (a Go duration).", + "type": "string" + } + } + }, "localai.FailoverChainsResponse": { "type": "object", "properties": { diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index 14224a6b499d..9c66faf4026c 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -1,5 +1,194 @@ basePath: / definitions: + cluster.Blocker: + properties: + forceable: + description: Forceable is true when a forced request goes ahead anyway. + type: boolean + id: + type: string + kind: + type: string + reason: + type: string + type: object + cluster.Carrier: + enum: + - nats + - tunnel + type: string + x-enum-varnames: + - CarrierNATS + - CarrierTunnel + cluster.CarrierRow: + properties: + active: + $ref: '#/definitions/cluster.Carrier' + changed_at: + description: |- + ChangedAt is stamped by the database on every change, and the age of a + change is read on that same clock (CarrierStore.DBNow), so that the + timeouts of a change need no clock that the replicas agree on. + type: string + changed_by: + type: string + draining: + $ref: '#/definitions/cluster.Carrier' + draining_until: + type: string + epoch: + type: integer + force: + description: |- + Force records that the admin who started the change accepted that a + replica which is not ready, or a worker that cannot follow, is left + behind. It belongs to one change and is cleared by the next one. + type: boolean + nats_url: + description: |- + NATSURL is the address of the NATS server that this change, or the last + change to NATS, builds from. The admin starts a change and the replicas + build in it, and an edit of the setting in between must not give two + replicas two servers, so every replica builds from this copy while a change + is under way. It never carries a credential. Empty when the cluster has no + NATS address. + type: string + note: + description: |- + Note says in a sentence how the last change ended, or why it did. It is + for operators. + type: string + prev_epoch: + type: integer + state: + $ref: '#/definitions/cluster.State' + target: + $ref: '#/definitions/cluster.Carrier' + type: object + cluster.InFlight: + properties: + claimed_claims: + type: integer + jobs: + description: Jobs is the number of jobs that run. + type: integer + loads: + description: Loads is the number of models that are loading. + type: integer + pending_claims: + description: PendingClaims and ClaimedClaims count the rows of the claim queue. + type: integer + type: object + cluster.ReplicaStatus: + properties: + availability: + additionalProperties: + type: string + description: |- + Availability maps each carrier to the reason the replica cannot build it, + empty when it can. A carrier it never looked at is absent. + type: object + availability_age_ns: + type: integer + id: + type: string + ready_epoch: + type: integer + ready_reason: + type: string + version: + type: string + type: object + cluster.Report: + properties: + active: + $ref: '#/definitions/cluster.Carrier' + blockers: + items: + $ref: '#/definitions/cluster.Blocker' + type: array + drain_remaining_ns: + description: |- + DrainRemaining is how long the previous carrier stays attached, when one + is draining. + type: integer + epoch: + type: integer + in_flight: + $ref: '#/definitions/cluster.InFlight' + ok: + type: boolean + replicas: + items: + $ref: '#/definitions/cluster.ReplicaStatus' + type: array + row: + $ref: '#/definitions/cluster.CarrierRow' + state: + $ref: '#/definitions/cluster.State' + target: + $ref: '#/definitions/cluster.Carrier' + timings: + $ref: '#/definitions/cluster.Timings' + warnings: + items: + type: string + type: array + workers: + items: + $ref: '#/definitions/cluster.WorkerStatus' + type: array + type: object + cluster.State: + enum: + - stable + - prepare + - commit + type: string + x-enum-varnames: + - StateStable + - StatePrepare + - StateCommit + cluster.Timings: + properties: + maxDrain: + description: |- + MaxDrain is how long the previous carrier stays attached after the commit. + Work that is still running on it then is failed by the reaper. + type: integer + prepareTimeout: + description: PrepareTimeout is how long prepare waits for every live replica. + type: integer + transitionWindow: + description: |- + TransitionWindow is how long commit waits for every live replica to + confirm. + type: integer + type: object + cluster.WorkerStatus: + properties: + attached: + items: + $ref: '#/definitions/cluster.Carrier' + type: array + can_follow: + description: CanFollow and Reason are set when the report has a target. + type: boolean + follow: + items: + $ref: '#/definitions/cluster.Carrier' + type: array + follow_error: + type: string + id: + type: string + name: + type: string + reason: + type: string + type: + type: string + type: object config.Gallery: properties: artifact_verification: @@ -460,6 +649,94 @@ definitions: logo_url: type: string type: object + localai.CarrierStatusResponse: + properties: + active: + $ref: '#/definitions/cluster.Carrier' + blockers: + items: + $ref: '#/definitions/cluster.Blocker' + type: array + drain_remaining_ns: + description: |- + DrainRemaining is how long the previous carrier stays attached, when one + is draining. + type: integer + epoch: + type: integer + in_flight: + $ref: '#/definitions/cluster.InFlight' + ok: + type: boolean + replicas: + items: + $ref: '#/definitions/cluster.ReplicaStatus' + type: array + row: + $ref: '#/definitions/cluster.CarrierRow' + settings: + $ref: '#/definitions/localai.ClusterSettingsView' + state: + $ref: '#/definitions/cluster.State' + target: + $ref: '#/definitions/cluster.Carrier' + timings: + $ref: '#/definitions/cluster.Timings' + warnings: + items: + type: string + type: array + workers: + items: + $ref: '#/definitions/cluster.WorkerStatus' + type: array + type: object + localai.CarrierSwitchRequest: + properties: + abort: + type: boolean + dry_run: + type: boolean + force: + type: boolean + target: + type: string + type: object + localai.ClusterSettingsRequest: + properties: + max_drain: + type: string + nats_url: + type: string + nats_worker_url: + type: string + prepare_timeout: + type: string + transition_window: + type: string + type: object + localai.ClusterSettingsView: + properties: + max_drain: + description: MaxDrain is how long the previous carrier stays attached after + a commit (a Go duration). + type: string + nats_url: + description: NATSURL is the NATS address that the frontends use. + type: string + nats_worker_url: + description: NATSWorkerURL is the NATS address that workers are told to use. + Empty means NATSURL. + type: string + prepare_timeout: + description: PrepareTimeout is how long a change waits for every replica to + be ready (a Go duration). + type: string + transition_window: + description: TransitionWindow is how long a replica may take to confirm a + commit (a Go duration). + type: string + type: object localai.FailoverChainsResponse: properties: chains: @@ -4144,6 +4421,141 @@ paths: summary: Upload a branding asset tags: - branding + /api/cluster/carrier: + get: + description: Returns the active carrier (nats or tunnel), the epoch, the state + of a change (stable, prepare or commit), the live replicas with their readiness, + the workers and whether each could follow a change to the other carrier, the + work in flight, and the cluster settings. Admin only. Answers 503 when distributed + mode is off. + responses: + "200": + description: OK + schema: + $ref: '#/definitions/localai.CarrierStatusResponse' + "503": + description: Distributed mode is not enabled + schema: + additionalProperties: + type: string + type: object + summary: Report the transport of a distributed cluster + tags: + - Nodes + post: + description: Takes a target (nats or tunnel) with dry_run and force, or abort. + A dry run returns the preflight report and changes nothing. A request starts + the change and returns 202 with the epoch to watch. A request that a blocker + stops returns 422 with the blockers; force accepts the blockers that are forceable. + A change already under way returns 409, and only abort is accepted then. Admin + only. + parameters: + - description: target, dry_run, force or abort + in: body + name: request + required: true + schema: + $ref: '#/definitions/localai.CarrierSwitchRequest' + responses: + "200": + description: Dry run report, or the row after an abort + schema: + $ref: '#/definitions/cluster.Report' + "202": + description: 'Change started: state, active, target, epoch, force, warnings, + row' + schema: + additionalProperties: true + type: object + "400": + description: Bad body or target + schema: + additionalProperties: + type: string + type: object + "409": + description: A change is under way, or there is nothing to abort + schema: + additionalProperties: + type: string + type: object + "422": + description: 'Blocked: error, blockers, report' + schema: + additionalProperties: true + type: object + "503": + description: Distributed mode is not enabled + schema: + additionalProperties: + type: string + type: object + summary: Dry-run, start or abort a change of carrier + tags: + - Nodes + /api/cluster/settings: + get: + description: Returns the NATS address of the frontends and of the workers and + the waits of a change. A user name, password or query string in an address + is dropped. Admin only. + responses: + "200": + description: OK + schema: + $ref: '#/definitions/localai.ClusterSettingsView' + "503": + description: Distributed mode is not enabled + schema: + additionalProperties: + type: string + type: object + summary: Read the settings of the cluster + tags: + - Nodes + put: + description: Stores the NATS address, the address workers use, or the waits + of a change. Saving a NATS address checks that the serving replica reaches + it and makes NATS available; it does not switch the cluster. Admin only. + parameters: + - description: settings to store + in: body + name: request + required: true + schema: + $ref: '#/definitions/localai.ClusterSettingsRequest' + responses: + "200": + description: saved, nats.reachable when a NATS address was checked, settings + schema: + additionalProperties: true + type: object + "400": + description: Unknown setting, bad value, or a credential in an address + schema: + additionalProperties: + type: string + type: object + "409": + description: A change of carrier is under way + schema: + additionalProperties: + type: string + type: object + "422": + description: This replica cannot reach the NATS server + schema: + additionalProperties: + type: string + type: object + "503": + description: Distributed mode is not enabled + schema: + additionalProperties: + type: string + type: object + summary: Store settings of the cluster + tags: + - Nodes /api/failover: get: produces: From 81f09466ebe4710fd4476f3b834789c26f580d5d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:55:15 +0000 Subject: [PATCH 149/162] feat(cli): add local-ai cluster to read and change the carrier The new command group calls the admin API of a running frontend: local-ai cluster carrier status local-ai cluster carrier switch --to nats|tunnel [--dry-run] [--force] [--yes] [--wait] local-ai cluster carrier abort local-ai cluster settings get local-ai cluster settings set [--nats-url ...] [--nats-worker-url ...] [--max-drain ...] A change runs the dry run first and prints the preflight with the list of live replicas. It then asks whether the list is complete and whether all frontends run this release, because a frontend of an older release cannot be detected. A command with no terminal needs --yes. --force is needed for blockers, and --wait follows the change until it commits or is aborted. Saving a NATS address does not change the carrier, and the command says so. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/cli/cli.go | 2 + core/cli/cluster/carrier.go | 300 +++++++++++++++++ core/cli/cluster/cluster.go | 148 +++++++++ core/cli/cluster/cluster_test.go | 416 ++++++++++++++++++++++++ core/cli/cluster/settings.go | 124 +++++++ docs/content/reference/cli-reference.md | 1 + 6 files changed, 991 insertions(+) create mode 100644 core/cli/cluster/carrier.go create mode 100644 core/cli/cluster/cluster.go create mode 100644 core/cli/cluster/cluster_test.go create mode 100644 core/cli/cluster/settings.go diff --git a/core/cli/cli.go b/core/cli/cli.go index 2b17ff96d5d0..b1ded4bda482 100644 --- a/core/cli/cli.go +++ b/core/cli/cli.go @@ -2,6 +2,7 @@ package cli import ( "github.com/mudler/LocalAI/core/cli/benchmark" + "github.com/mudler/LocalAI/core/cli/cluster" cliContext "github.com/mudler/LocalAI/core/cli/context" "github.com/mudler/LocalAI/core/cli/worker" ) @@ -10,6 +11,7 @@ var CLI struct { cliContext.Context `embed:""` Benchmark benchmark.Command `cmd:"" help:"Benchmark configured text models against a running LocalAI server"` + Cluster cluster.Command `cmd:"" help:"Read or change the carrier and the settings of a distributed cluster through a running frontend"` Run RunCMD `cmd:"" help:"Run LocalAI, this the default command if no other command is specified. Run 'local-ai run --help' for more information" default:"withargs"` Chat ChatCMD `cmd:"" help:"Run the built-in terminal agent against a LocalAI server"` Federated FederatedCLI `cmd:"" help:"Run LocalAI in federated mode"` diff --git a/core/cli/cluster/carrier.go b/core/cli/cluster/carrier.go new file mode 100644 index 000000000000..565dc8b158d7 --- /dev/null +++ b/core/cli/cluster/carrier.go @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "strings" + "text/tabwriter" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +const carrierPath = "/api/cluster/carrier" + +// CarrierCommand groups the commands about the carrier. +type CarrierCommand struct { + Status StatusCommand `cmd:"" default:"withargs" help:"Show the active carrier, the state of a change, the replicas and the workers"` + Switch SwitchCommand `cmd:"" help:"Change the carrier of the cluster, or check what would block the change (--dry-run)"` + Abort AbortCommand `cmd:"" help:"Abort a change that is being prepared"` +} + +// StatusCommand shows the state of the carrier. +type StatusCommand struct { + Remote `embed:""` +} + +// SwitchCommand changes the carrier. +type SwitchCommand struct { + To string `required:"" enum:"nats,tunnel" help:"Carrier to change to [${enum}]."` + DryRun bool `name:"dry-run" help:"Run the preflight and change nothing. It lists what would block the change."` + Force bool `help:"Go ahead although some workers cannot follow or some replicas are not ready. Those workers become unroutable until they are fixed."` + Yes bool `help:"Confirm that the replicas that are listed are all the frontends of the cluster, and that all run this release. Without it the command asks, and a command with no terminal refuses."` + Wait bool `help:"Wait until the change is committed, and print the progress."` + // WaitTimeout bounds --wait. + WaitTimeout time.Duration `name:"wait-timeout" default:"10m" help:"How long --wait waits."` + Remote `embed:""` +} + +// AbortCommand aborts a change in the prepare state. +type AbortCommand struct { + Remote `embed:""` +} + +// pollInterval is how often --wait reads the state. +var pollInterval = 2 * time.Second + +func (c *StatusCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *StatusCommand) run(ctx context.Context, out io.Writer) error { + var view cluster.Report + _, raw, err := c.do(ctx, c.client(), http.MethodGet, carrierPath, nil, &view) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + printStatus(out, view) + return nil +} + +func (c *AbortCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *AbortCommand) run(ctx context.Context, out io.Writer) error { + var row cluster.CarrierRow + _, raw, err := c.do(ctx, c.client(), http.MethodPost, carrierPath, map[string]any{"abort": true}, &row) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + _, _ = fmt.Fprintf(out, "Aborted. The cluster stays on %s (epoch %d).\n", row.Active, row.Epoch) + return nil +} + +func (c *SwitchCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout, os.Stdin, stdinIsTerminal()) +} + +func stdinIsTerminal() bool { + fi, err := os.Stdin.Stat() + return err == nil && fi.Mode()&os.ModeCharDevice != 0 +} + +func (c *SwitchCommand) run(ctx context.Context, out io.Writer, in io.Reader, interactive bool) error { + client := c.client() + target := c.To + + // The dry run is the preflight. A change runs it again on the server, so + // this one is for the person who reads the replica list. + var report cluster.Report + _, raw, err := c.do(ctx, client, http.MethodPost, carrierPath, map[string]any{"target": target, "dry_run": true}, &report) + if err != nil { + return err + } + if c.DryRun { + if c.JSON { + return printJSON(out, raw) + } + printPreflight(out, report, target) + if !report.OK { + return errors.New("the change would be blocked; fix the blockers, or use --force for those that can be forced") + } + return nil + } + + if !c.JSON { + printPreflight(out, report, target) + } + if !report.OK && !c.Force { + return errors.New("the change is blocked; fix the blockers, or use --force for those that can be forced") + } + if err := c.confirm(out, in, interactive, report); err != nil { + return err + } + + var started struct { + State string `json:"state"` + Active string `json:"active"` + Target string `json:"target"` + Epoch int64 `json:"epoch"` + } + _, raw, err = c.do(ctx, client, http.MethodPost, carrierPath, map[string]any{"target": target, "force": c.Force}, &started) + if err != nil { + var ae *apiError + if errors.As(err, &ae) && ae.Status == http.StatusUnprocessableEntity { + var body struct { + Blockers []cluster.Blocker `json:"blockers"` + } + _ = json.Unmarshal(ae.Body, &body) + var b strings.Builder + for _, bl := range body.Blockers { + fmt.Fprintf(&b, "\n %s %s: %s", bl.Kind, bl.ID, bl.Reason) + } + return fmt.Errorf("the change is blocked:%s", b.String()) + } + return err + } + if c.JSON && !c.Wait { + return printJSON(out, raw) + } + _, _ = fmt.Fprintf(out, "Change to %s started (epoch %d, state %s).\n", started.Target, started.Epoch, started.State) + if !c.Wait { + _, _ = fmt.Fprintln(out, "The replicas carry it out. Run `local-ai cluster carrier status` to follow it.") + return nil + } + return c.wait(ctx, client, out, started.Epoch) +} + +// confirm asks the person to confirm the replica list, which the server cannot +// check: a frontend of an older release writes no row and is not in the list. +func (c *SwitchCommand) confirm(out io.Writer, in io.Reader, interactive bool, report cluster.Report) error { + if c.Yes { + return nil + } + if !interactive { + return errors.New("confirm that the replicas listed above are all the frontends of the cluster and that all run this release: pass --yes") + } + q := fmt.Sprintf("Are these %d replicas all the frontends of the cluster, and do all of them run this release?", len(report.Replicas)) + if c.Force { + q += " --force makes the workers that cannot follow unroutable." + } + _, _ = fmt.Fprintf(out, "%s [y/N] ", q) + line, _ := bufio.NewReader(in).ReadString('\n') + switch strings.ToLower(strings.TrimSpace(line)) { + case "y", "yes": + return nil + } + return errors.New("not confirmed; nothing was changed") +} + +// wait reads the state until the change is committed or gone. +func (c *SwitchCommand) wait(ctx context.Context, client *http.Client, out io.Writer, epoch int64) error { + ctx, cancel := context.WithTimeout(ctx, c.WaitTimeout) + defer cancel() + last := "" + for { + var st cluster.Report + if _, _, err := c.do(ctx, client, http.MethodGet, carrierPath, nil, &st); err != nil { + if ctx.Err() != nil { + return fmt.Errorf("gave up waiting; the change goes on, check it with `local-ai cluster carrier status`: %w", ctx.Err()) + } + return err + } + line := fmt.Sprintf("state %s, active %s, epoch %d", st.State, st.Active, st.Epoch) + if line != last { + _, _ = fmt.Fprintln(out, line) + last = line + } + switch { + case st.State == cluster.StateStable && string(st.Active) == c.To: + _, _ = fmt.Fprintf(out, "The cluster is on %s.\n", c.To) + if st.DrainRemaining > 0 { + _, _ = fmt.Fprintf(out, "The previous carrier drains for %s more.\n", st.DrainRemaining.Round(time.Second)) + } + return nil + case st.State == cluster.StateStable && st.Epoch > epoch: + note := st.Row.Note + if note != "" { + note = ": " + note + } + return fmt.Errorf("the change was aborted; the cluster stays on %s%s", st.Active, note) + } + select { + case <-ctx.Done(): + return fmt.Errorf("gave up waiting; the change goes on, check it with `local-ai cluster carrier status`: %w", ctx.Err()) + case <-time.After(pollInterval): + } + } +} + +func printStatus(out io.Writer, r cluster.Report) { + _, _ = fmt.Fprintf(out, "Active carrier: %s (epoch %d)\n", r.Active, r.Epoch) + _, _ = fmt.Fprintf(out, "State: %s", r.State) + if r.Target != "" { + _, _ = fmt.Fprintf(out, " (changing to %s)", r.Target) + } + _, _ = fmt.Fprintln(out) + if r.Row.Draining != "" { + _, _ = fmt.Fprintf(out, "Draining: %s", r.Row.Draining) + if r.DrainRemaining > 0 { + _, _ = fmt.Fprintf(out, ", %s left", r.DrainRemaining.Round(time.Second)) + } + _, _ = fmt.Fprintln(out) + } + printReplicas(out, r) + printWorkers(out, r) +} + +func printReplicas(out io.Writer, r cluster.Report) { + _, _ = fmt.Fprintf(out, "\nReplicas (%d):\n", len(r.Replicas)) + tw := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(tw, " ID\tVERSION\tREADY EPOCH\tNOTE") + for _, rep := range r.Replicas { + _, _ = fmt.Fprintf(tw, " %s\t%s\t%d\t%s\n", rep.ID, rep.Version, rep.ReadyEpoch, rep.ReadyReason) + } + _ = tw.Flush() +} + +func printWorkers(out io.Writer, r cluster.Report) { + _, _ = fmt.Fprintf(out, "\nWorkers (%d):\n", len(r.Workers)) + tw := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(tw, " NAME\tID\tATTACHED\tCAN FOLLOW\tNOTE") + for _, w := range r.Workers { + var att []string + for _, a := range w.Attached { + att = append(att, string(a)) + } + note := w.Reason + if w.FollowError != "" { + if note != "" { + note += "; " + } + note += "follow_error: " + w.FollowError + } + _, _ = fmt.Fprintf(tw, " %s\t%s\t%s\t%t\t%s\n", w.Name, w.ID, strings.Join(att, ","), w.CanFollow, note) + } + _ = tw.Flush() +} + +func printPreflight(out io.Writer, r cluster.Report, target string) { + _, _ = fmt.Fprintf(out, "Preflight for a change from %s to %s:\n", r.Active, target) + if r.OK { + _, _ = fmt.Fprintln(out, " Nothing blocks the change.") + } + for _, b := range r.Blockers { + force := "" + if b.Forceable { + force = " (can be forced)" + } + _, _ = fmt.Fprintf(out, " BLOCKER %s %s: %s%s\n", b.Kind, b.ID, b.Reason, force) + } + for _, w := range r.Warnings { + _, _ = fmt.Fprintf(out, " warning: %s\n", w) + } + _, _ = fmt.Fprintf(out, " In flight: %d loads, %d jobs, %d pending claims, %d claimed\n", + r.InFlight.Loads, r.InFlight.Jobs, r.InFlight.PendingClaims, r.InFlight.ClaimedClaims) + printReplicas(out, r) + _, _ = fmt.Fprintln(out, " Frontends of an older release write no row and are not listed. Check that this list is complete.") + printWorkers(out, r) +} diff --git a/core/cli/cluster/cluster.go b/core/cli/cluster/cluster.go new file mode 100644 index 000000000000..1c50e6c898af --- /dev/null +++ b/core/cli/cluster/cluster.go @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: MIT + +// Package cluster is the "local-ai cluster" command. It calls the admin API of +// a running distributed frontend to read or change the carrier of the cluster +// and to read or store the cluster settings. +package cluster + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" +) + +// Command is the "local-ai cluster" command group. +type Command struct { + Carrier CarrierCommand `cmd:"" help:"Read the carrier of a distributed cluster, or change it"` + Settings SettingsCommand `cmd:"" help:"Read or store the settings of a distributed cluster"` +} + +// Remote says which server a command talks to. +type Remote struct { + Endpoint string `default:"http://127.0.0.1:8080" env:"LOCALAI_ENDPOINT" help:"URL of a LocalAI frontend of the cluster."` + APIKey string `name:"api-key" env:"LOCALAI_API_KEY,API_KEY" help:"API key of an admin user."` + Timeout time.Duration `default:"2m" help:"Timeout for each request. A dry run asks every replica, so it can take a few seconds."` + JSON bool `name:"json" help:"Write the answer of the server as JSON."` +} + +// apiError is an answer that is not a success. +type apiError struct { + Status int + Msg string + // Body is the decoded body, for the callers that read more than the message. + Body []byte +} + +func (e *apiError) Error() string { + switch e.Status { + case http.StatusUnauthorized, http.StatusForbidden: + return fmt.Sprintf("the server refused the request (HTTP %d): use --api-key with the key of an admin user", e.Status) + case http.StatusServiceUnavailable: + return "the server is not running in distributed mode, or the cluster is not ready: " + e.Msg + } + if e.Msg == "" { + return fmt.Sprintf("HTTP status %d", e.Status) + } + return fmt.Sprintf("%s (HTTP %d)", e.Msg, e.Status) +} + +// endpointURL joins the base URL and the path of an API route. +func endpointURL(base, path string) (string, error) { + u, err := url.Parse(base) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(base, "#") { + return "", errors.New("endpoint must be an HTTP(S) URL without credentials, query, or fragment") + } + u.Path = strings.TrimRight(u.Path, "/") + path + u.RawPath = "" + return u.String(), nil +} + +// do sends one request and decodes a JSON answer into out. It returns the status +// code and the raw body as well. A status outside 2xx is an *apiError. +func (r *Remote) do(ctx context.Context, client *http.Client, method, path string, in, out any) (int, []byte, error) { + target, err := endpointURL(r.Endpoint, path) + if err != nil { + return 0, nil, err + } + var body io.Reader + if in != nil { + raw, err := json.Marshal(in) + if err != nil { + return 0, nil, err + } + body = bytes.NewReader(raw) + } + req, err := http.NewRequestWithContext(ctx, method, target, body) + if err != nil { + return 0, nil, errors.New("cannot create the request") + } + req.Header.Set("Accept", "application/json") + if in != nil { + req.Header.Set("Content-Type", "application/json") + } + if r.APIKey != "" { + req.Header.Set("Authorization", "Bearer "+r.APIKey) + } + resp, err := client.Do(req) + if err != nil { + if ctx.Err() != nil { + return 0, nil, ctx.Err() + } + // A transport error can echo the URL and the credentials of a proxy. + return 0, nil, errors.New("cannot reach the server at the endpoint") + } + defer func() { _ = resp.Body.Close() }() + raw, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + if err != nil { + return resp.StatusCode, nil, errors.New("cannot read the answer") + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + var e struct { + Error string `json:"error"` + } + _ = json.Unmarshal(raw, &e) + return resp.StatusCode, raw, &apiError{Status: resp.StatusCode, Msg: e.Error, Body: raw} + } + if out != nil { + if err := json.Unmarshal(raw, out); err != nil { + return resp.StatusCode, raw, fmt.Errorf("the server answered with a body that is not valid: %w", err) + } + } + return resp.StatusCode, raw, nil +} + +func (r *Remote) client() *http.Client { + timeout := r.Timeout + if timeout <= 0 { + timeout = 2 * time.Minute + } + return httpclient.NewWithTimeout(timeout) +} + +// printJSON writes raw, indented. +func printJSON(w io.Writer, raw []byte) error { + var buf bytes.Buffer + if err := json.Indent(&buf, raw, "", " "); err != nil { + _, err = w.Write(raw) + return err + } + buf.WriteByte('\n') + _, err := w.Write(buf.Bytes()) + return err +} + +func signalContext() (context.Context, context.CancelFunc) { + return signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) +} diff --git a/core/cli/cluster/cluster_test.go b/core/cli/cluster/cluster_test.go new file mode 100644 index 000000000000..416d3f867ff3 --- /dev/null +++ b/core/cli/cluster/cluster_test.go @@ -0,0 +1,416 @@ +// SPDX-License-Identifier: MIT +package cluster + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/alecthomas/kong" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestClusterCLI(t *testing.T) { RegisterFailHandler(Fail); RunSpecs(t, "Cluster CLI") } + +// fakeServer records the requests of a command and answers from a table. +type fakeServer struct { + mu sync.Mutex + requests []recorded + handler func(r *http.Request, body map[string]any) (int, any) + srv *httptest.Server +} + +type recorded struct { + Method, Path, Auth string + Body map[string]any +} + +func newFakeServer(h func(r *http.Request, body map[string]any) (int, any)) *fakeServer { + f := &fakeServer{handler: h} + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + f.mu.Lock() + f.requests = append(f.requests, recorded{r.Method, r.URL.Path, r.Header.Get("Authorization"), body}) + f.mu.Unlock() + code, resp := f.handler(r, body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + _ = json.NewEncoder(w).Encode(resp) + })) + return f +} + +func (f *fakeServer) last() recorded { + f.mu.Lock() + defer f.mu.Unlock() + return f.requests[len(f.requests)-1] +} + +func (f *fakeServer) count() int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.requests) +} + +func remote(f *fakeServer) Remote { + return Remote{Endpoint: f.srv.URL, APIKey: "k", Timeout: 5 * time.Second} +} + +var preflightOK = map[string]any{ + "active": "nats", "epoch": 3, "state": "stable", "ok": true, + "replicas": []map[string]any{{"id": "r1", "version": "v1", "ready_epoch": 3}, {"id": "r2", "version": "v1", "ready_epoch": 3}}, + "workers": []map[string]any{{"id": "w1", "name": "gpu", "attached": []string{"nats"}, "can_follow": true}}, + "in_flight": map[string]any{"loads": 1}, +} + +var preflightBlocked = map[string]any{ + "active": "nats", "epoch": 3, "state": "stable", "ok": false, + "blockers": []map[string]any{{"kind": "worker", "id": "w2", "reason": "no routable address", "forceable": true}}, + "replicas": []map[string]any{{"id": "r1", "version": "v1", "ready_epoch": 3}}, + "workers": []map[string]any{{"id": "w2", "name": "edge", "attached": []string{"nats"}, "can_follow": false, "reason": "no routable address", "follow_error": "needs address"}}, +} + +var _ = Describe("local-ai cluster", func() { + var out bytes.Buffer + BeforeEach(func() { out.Reset() }) + + Describe("parsing", func() { + It("reads the subcommands and their flags", func() { + var c struct { + Cluster Command `cmd:""` + } + parser, err := kong.New(&c) + Expect(err).NotTo(HaveOccurred()) + + ctx, err := parser.Parse([]string{"cluster", "carrier", "switch", "--to", "tunnel", "--dry-run", "--force", "--yes", "--wait"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ctx.Command()).To(Equal("cluster carrier switch")) + sw := c.Cluster.Carrier.Switch + Expect(sw.To).To(Equal("tunnel")) + Expect(sw.DryRun && sw.Force && sw.Yes && sw.Wait).To(BeTrue()) + Expect(sw.Endpoint).To(Equal("http://127.0.0.1:8080")) + + _, err = parser.Parse([]string{"cluster", "carrier", "switch", "--to", "kafka"}) + Expect(err).To(HaveOccurred()) + var bare struct { + Cluster Command `cmd:""` + } + bareParser, err := kong.New(&bare) + Expect(err).NotTo(HaveOccurred()) + _, err = bareParser.Parse([]string{"cluster", "carrier", "switch"}) + Expect(err).To(HaveOccurred()) + + ctx, err = parser.Parse([]string{"cluster", "carrier", "abort"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ctx.Command()).To(Equal("cluster carrier abort")) + + ctx, err = parser.Parse([]string{"cluster", "carrier"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ctx.Command()).To(Equal("cluster carrier status")) + }) + + It("tells an absent setting from an empty one", func() { + var c struct { + Cluster Command `cmd:""` + } + parser, err := kong.New(&c) + Expect(err).NotTo(HaveOccurred()) + _, err = parser.Parse([]string{"cluster", "settings", "set", "--nats-url=", "--max-drain", "20m"}) + Expect(err).NotTo(HaveOccurred()) + set := c.Cluster.Settings.Set + Expect(set.NATSURL).NotTo(BeNil()) + Expect(*set.NATSURL).To(BeEmpty()) + Expect(set.NATSWorkerURL).To(BeNil()) + Expect(*set.MaxDrain).To(Equal("20m")) + }) + }) + + Describe("carrier status", func() { + It("prints the carrier, the replicas and the workers, and sends the key", func() { + f := newFakeServer(func(r *http.Request, _ map[string]any) (int, any) { + return 200, map[string]any{ + "active": "tunnel", "epoch": 9, "state": "stable", + "row": map[string]any{"draining": "nats"}, + "drain_remaining_ns": int64(90 * time.Second), + "replicas": preflightOK["replicas"], "workers": preflightBlocked["workers"], + } + }) + defer f.srv.Close() + c := &StatusCommand{Remote: remote(f)} + Expect(c.run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Method).To(Equal("GET")) + Expect(f.last().Path).To(Equal("/api/cluster/carrier")) + Expect(f.last().Auth).To(Equal("Bearer k")) + s := out.String() + Expect(s).To(ContainSubstring("Active carrier: tunnel (epoch 9)")) + Expect(s).To(ContainSubstring("Draining: nats, 1m30s left")) + Expect(s).To(ContainSubstring("Replicas (2)")) + Expect(s).To(ContainSubstring("follow_error: needs address")) + }) + + It("writes the JSON of the server with --json", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + r := remote(f) + r.JSON = true + Expect((&StatusCommand{Remote: r}).run(context.Background(), &out)).To(Succeed()) + var got map[string]any + Expect(json.Unmarshal(out.Bytes(), &got)).To(Succeed()) + Expect(got["active"]).To(Equal("nats")) + }) + + It("explains a refused key and a frontend that is not distributed", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 401, map[string]any{"error": "no"} }) + err := (&StatusCommand{Remote: remote(f)}).run(context.Background(), &out) + f.srv.Close() + Expect(err).To(MatchError(ContainSubstring("--api-key"))) + + f = newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 503, map[string]any{"error": "distributed mode is not enabled on this frontend"} + }) + defer f.srv.Close() + err = (&StatusCommand{Remote: remote(f)}).run(context.Background(), &out) + Expect(err).To(MatchError(ContainSubstring("not running in distributed mode"))) + }) + + It("refuses an endpoint that is not a plain HTTP URL", func() { + err := (&StatusCommand{Remote: Remote{Endpoint: "http://u:p@host:1"}}).run(context.Background(), &out) + Expect(err).To(MatchError(ContainSubstring("without credentials"))) + }) + }) + + Describe("carrier switch", func() { + It("only runs the preflight with --dry-run", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", DryRun: true, Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader(""), false)).To(Succeed()) + Expect(f.count()).To(Equal(1)) + Expect(f.last().Body).To(HaveKeyWithValue("dry_run", true)) + Expect(f.last().Body).To(HaveKeyWithValue("target", "tunnel")) + Expect(out.String()).To(ContainSubstring("Nothing blocks the change.")) + Expect(out.String()).To(ContainSubstring("Check that this list is complete.")) + }) + + It("fails a dry run that has a blocker and lists it", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightBlocked }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", DryRun: true, Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("blocked"))) + Expect(out.String()).To(ContainSubstring("BLOCKER worker w2: no routable address (can be forced)")) + }) + + It("refuses to start a change without a confirmation when there is no terminal", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("--yes"))) + Expect(f.count()).To(Equal(1), "only the dry run was sent") + }) + + It("does not start a change when the answer to the question is no", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader("n\n"), true) + Expect(err).To(MatchError(ContainSubstring("nothing was changed"))) + Expect(f.count()).To(Equal(1)) + Expect(out.String()).To(ContainSubstring("Are these 2 replicas all the frontends")) + }) + + It("starts the change after a yes at the question", func() { + f := newFakeServer(func(_ *http.Request, b map[string]any) (int, any) { + if b["dry_run"] == true { + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader("yes\n"), true)).To(Succeed()) + Expect(f.last().Body).To(HaveKeyWithValue("target", "tunnel")) + Expect(f.last().Body).To(HaveKeyWithValue("force", false)) + Expect(f.last().Body).NotTo(HaveKey("dry_run")) + Expect(out.String()).To(ContainSubstring("Change to tunnel started (epoch 4, state prepare).")) + }) + + It("stops at a blocker unless --force, and sends force when given", func() { + var forced bool + f := newFakeServer(func(_ *http.Request, b map[string]any) (int, any) { + if b["dry_run"] == true { + return 200, preflightBlocked + } + forced = b["force"] == true + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + err := (&SwitchCommand{To: "tunnel", Yes: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("blocked"))) + Expect(f.count()).To(Equal(1)) + + Expect((&SwitchCommand{To: "tunnel", Yes: true, Force: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false)).To(Succeed()) + Expect(forced).To(BeTrue()) + }) + + It("shows the blockers of a 422 and the busy state of a 409", func() { + f := newFakeServer(func(_ *http.Request, b map[string]any) (int, any) { + if b["dry_run"] == true { + return 200, preflightOK + } + return 422, map[string]any{"error": "blocked", "blockers": []map[string]any{{"kind": "replica", "id": "r9", "reason": "not ready"}}} + }) + err := (&SwitchCommand{To: "tunnel", Yes: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false) + f.srv.Close() + Expect(err).To(MatchError(ContainSubstring("replica r9: not ready"))) + + f = newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 409, map[string]any{"error": "a change is under way"} + }) + defer f.srv.Close() + err = (&SwitchCommand{To: "tunnel", DryRun: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("a change is under way"))) + }) + + It("waits for the commit with --wait", func() { + old := pollInterval + pollInterval = time.Millisecond + DeferCleanup(func() { pollInterval = old }) + var polls int + f := newFakeServer(func(r *http.Request, b map[string]any) (int, any) { + switch { + case r.Method == http.MethodGet: + polls++ + if polls < 3 { + return 200, map[string]any{"active": "nats", "state": "prepare", "epoch": 4} + } + return 200, map[string]any{"active": "tunnel", "state": "stable", "epoch": 6, "drain_remaining_ns": int64(15 * time.Minute)} + case b["dry_run"] == true: + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Yes: true, Wait: true, WaitTimeout: 5 * time.Second, Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader(""), false)).To(Succeed()) + Expect(out.String()).To(ContainSubstring("The cluster is on tunnel.")) + Expect(out.String()).To(ContainSubstring("drains for 15m0s more")) + }) + + It("reports an abort that --wait sees", func() { + old := pollInterval + pollInterval = time.Millisecond + DeferCleanup(func() { pollInterval = old }) + f := newFakeServer(func(r *http.Request, b map[string]any) (int, any) { + switch { + case r.Method == http.MethodGet: + return 200, map[string]any{"active": "nats", "state": "stable", "epoch": 5, "row": map[string]any{"note": "replica r2 not ready"}} + case b["dry_run"] == true: + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Yes: true, Wait: true, WaitTimeout: 5 * time.Second, Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("aborted"))) + Expect(err).To(MatchError(ContainSubstring("replica r2 not ready"))) + }) + + It("gives up waiting when the time is out", func() { + old := pollInterval + pollInterval = time.Millisecond + DeferCleanup(func() { pollInterval = old }) + f := newFakeServer(func(r *http.Request, b map[string]any) (int, any) { + switch { + case r.Method == http.MethodGet: + return 200, map[string]any{"active": "nats", "state": "prepare", "epoch": 4} + case b["dry_run"] == true: + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Yes: true, Wait: true, WaitTimeout: 30 * time.Millisecond, Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader(""), false)).To(MatchError(ContainSubstring("gave up waiting"))) + }) + }) + + Describe("carrier abort", func() { + It("posts abort and says where the cluster stays", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"active": "nats", "epoch": 8, "state": "stable"} + }) + defer f.srv.Close() + Expect((&AbortCommand{Remote: remote(f)}).run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Body).To(HaveKeyWithValue("abort", true)) + Expect(out.String()).To(ContainSubstring("The cluster stays on nats (epoch 8).")) + }) + + It("passes on the 409 of a change that cannot be aborted", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 409, map[string]any{"error": "nothing to abort"} + }) + defer f.srv.Close() + Expect((&AbortCommand{Remote: remote(f)}).run(context.Background(), &out)).To(MatchError(ContainSubstring("nothing to abort"))) + }) + }) + + Describe("settings", func() { + It("shows the settings", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"nats_url": "nats://n:4222", "max_drain": "15m"} + }) + defer f.srv.Close() + Expect((&GetSettingsCommand{Remote: remote(f)}).run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Path).To(Equal("/api/cluster/settings")) + Expect(out.String()).To(ContainSubstring("nats_url: nats://n:4222")) + Expect(out.String()).To(ContainSubstring("nats_worker_url: (not set)")) + }) + + It("sends only the flags that are given, and says the carrier did not change", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"saved": true, "nats": map[string]any{"reachable": true}, "settings": map[string]any{"nats_url": "nats://n:4222"}} + }) + defer f.srv.Close() + url := "nats://n:4222" + c := &SetSettingsCommand{NATSURL: &url, Remote: remote(f)} + Expect(c.run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Method).To(Equal("PUT")) + Expect(f.last().Body).To(Equal(map[string]any{"nats_url": url})) + Expect(out.String()).To(ContainSubstring("reaches the NATS server")) + Expect(out.String()).To(ContainSubstring("The carrier did not change.")) + }) + + It("sends an empty value to clear a setting", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"saved": true, "settings": map[string]any{}} + }) + defer f.srv.Close() + empty := "" + Expect((&SetSettingsCommand{NATSWorkerURL: &empty, Remote: remote(f)}).run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Body).To(Equal(map[string]any{"nats_worker_url": ""})) + }) + + It("refuses a call with nothing to store, and shows a refused address", func() { + Expect((&SetSettingsCommand{}).run(context.Background(), io.Discard)).To(MatchError(ContainSubstring("at least one setting"))) + + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 422, map[string]any{"error": "nats_url: this replica cannot reach the NATS server at nats://x:1"} + }) + defer f.srv.Close() + u := "nats://x:1" + Expect((&SetSettingsCommand{NATSURL: &u, Remote: remote(f)}).run(context.Background(), &out)).To(MatchError(ContainSubstring("cannot reach the NATS server"))) + }) + }) +}) diff --git a/core/cli/cluster/settings.go b/core/cli/cluster/settings.go new file mode 100644 index 000000000000..6cfcd52807d6 --- /dev/null +++ b/core/cli/cluster/settings.go @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "os" +) + +const settingsPath = "/api/cluster/settings" + +// SettingsCommand groups the commands about the settings of the cluster. +type SettingsCommand struct { + Get GetSettingsCommand `cmd:"" default:"withargs" help:"Show the settings"` + Set SetSettingsCommand `cmd:"" help:"Store settings. Storing a NATS address does not change the carrier."` +} + +// GetSettingsCommand shows the settings. +type GetSettingsCommand struct { + Remote `embed:""` +} + +// SetSettingsCommand stores settings. A flag that is not given leaves the +// setting as it is, and a flag given an empty value clears it. +type SetSettingsCommand struct { + NATSURL *string `name:"nats-url" help:"NATS address that the frontends use. The serving replica checks that it reaches the server. Empty clears it."` + NATSWorkerURL *string `name:"nats-worker-url" help:"NATS address that workers are told to use, when it differs from --nats-url. Empty clears it."` + PrepareTimeout *string `name:"prepare-timeout" help:"How long a change waits for every replica to be ready, as a duration such as 60s."` + TransitionWindow *string `name:"transition-window" help:"How long a replica may take to confirm a commit, as a duration."` + MaxDrain *string `name:"max-drain" help:"How long the previous carrier stays attached after a commit, as a duration such as 15m."` + Remote `embed:""` +} + +type settingsView struct { + NATSURL string `json:"nats_url"` + NATSWorkerURL string `json:"nats_worker_url"` + PrepareTimeout string `json:"prepare_timeout"` + TransitionWindow string `json:"transition_window"` + MaxDrain string `json:"max_drain"` +} + +func (c *GetSettingsCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *GetSettingsCommand) run(ctx context.Context, out io.Writer) error { + var view settingsView + _, raw, err := c.do(ctx, c.client(), http.MethodGet, settingsPath, nil, &view) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + printSettings(out, view) + return nil +} + +func (c *SetSettingsCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *SetSettingsCommand) run(ctx context.Context, out io.Writer) error { + body := map[string]string{} + for name, v := range map[string]*string{ + "nats_url": c.NATSURL, + "nats_worker_url": c.NATSWorkerURL, + "prepare_timeout": c.PrepareTimeout, + "transition_window": c.TransitionWindow, + "max_drain": c.MaxDrain, + } { + if v != nil { + body[name] = *v + } + } + if len(body) == 0 { + return errors.New("give at least one setting to store, for example --nats-url nats://host:4222") + } + var answer struct { + Saved bool `json:"saved"` + NATS *struct { + Reachable bool `json:"reachable"` + } `json:"nats"` + Settings settingsView `json:"settings"` + } + _, raw, err := c.do(ctx, c.client(), http.MethodPut, settingsPath, body, &answer) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + _, _ = fmt.Fprintln(out, "Saved.") + if answer.NATS != nil && answer.NATS.Reachable { + _, _ = fmt.Fprintln(out, "The serving replica reaches the NATS server.") + } + if c.NATSURL != nil || c.NATSWorkerURL != nil { + _, _ = fmt.Fprintln(out, "The carrier did not change. To use NATS, run `local-ai cluster carrier switch --to nats --dry-run`, then the same command without --dry-run.") + } + printSettings(out, answer.Settings) + return nil +} + +func printSettings(out io.Writer, v settingsView) { + show := func(s string) string { + if s == "" { + return "(not set)" + } + return s + } + _, _ = fmt.Fprintf(out, "nats_url: %s\n", show(v.NATSURL)) + _, _ = fmt.Fprintf(out, "nats_worker_url: %s\n", show(v.NATSWorkerURL)) + _, _ = fmt.Fprintf(out, "prepare_timeout: %s\n", show(v.PrepareTimeout)) + _, _ = fmt.Fprintf(out, "transition_window: %s\n", show(v.TransitionWindow)) + _, _ = fmt.Fprintf(out, "max_drain: %s\n", show(v.MaxDrain)) +} diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 31bb79013d12..4cea07cb5d26 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -224,6 +224,7 @@ LocalAI supports several subcommands beyond `run`: - `local-ai agent` - Run agents standalone without the full LocalAI server - `local-ai mcp-server` - Run the LocalAI admin tool surface as a stdio MCP server (controls a remote LocalAI instance over HTTP) - `local-ai worker` - Start a worker for distributed mode (generic, backend-agnostic) +- `local-ai cluster` - Read or change the carrier (NATS or the database tunnel) and the settings of a distributed cluster through a running frontend. See [Distributed Mode]({{% relref "features/distributed-mode" %}}) - `local-ai p2p-worker` - Run workers to distribute workload via p2p (llama.cpp-only) - `local-ai agent-worker` - Start an agent worker for distributed mode (executes agent chats via NATS) - `local-ai util` - Utility commands From 81c1a8166df88c1bd4c533967ddd37e5788d9c40 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 8 Oct 2026 16:58:30 +0000 Subject: [PATCH 150/162] feat(ui): add a cluster transport panel to the Nodes page The panel shows the active carrier, the epoch, the state of a change (stable, prepare, commit, draining), the readiness of each frontend replica, and each worker with its profile, the carriers it is attached to and its follow error. It shows only when the cluster routes answer, so it does not appear on a frontend that is not distributed. The switch dialog runs a dry run and lists the blockers and the live replica list. The admin must confirm that the list is complete and that all frontends are upgraded. A forced change needs a second confirmation. An admin can abort a change during prepare, and a drain shows a countdown. The NATS settings form saves the address and reports whether this frontend reaches it. Saving does not switch the cluster. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- .../react-ui/e2e/cluster-transport.spec.js | 226 +++++++++++ core/http/react-ui/src/App.css | 32 ++ .../src/components/nodes/ClusterTransport.jsx | 364 ++++++++++++++++++ core/http/react-ui/src/pages/Nodes.jsx | 3 + core/http/react-ui/src/utils/api.js | 15 + .../react-ui/src/utils/clusterTransport.js | 66 ++++ .../src/utils/clusterTransport.test.js | 37 ++ core/http/react-ui/src/utils/config.js | 2 + 8 files changed, 745 insertions(+) create mode 100644 core/http/react-ui/e2e/cluster-transport.spec.js create mode 100644 core/http/react-ui/src/components/nodes/ClusterTransport.jsx create mode 100644 core/http/react-ui/src/utils/clusterTransport.js create mode 100644 core/http/react-ui/src/utils/clusterTransport.test.js diff --git a/core/http/react-ui/e2e/cluster-transport.spec.js b/core/http/react-ui/e2e/cluster-transport.spec.js new file mode 100644 index 000000000000..014195fec395 --- /dev/null +++ b/core/http/react-ui/e2e/cluster-transport.spec.js @@ -0,0 +1,226 @@ +import { test, expect } from './coverage-fixtures.js' + +const node = { id: 'n1', name: 'atlas', node_type: 'backend', address: '10.0.0.1:50051', status: 'healthy', labels: {}, total_vram: 100, available_vram: 40, total_ram: 200, available_ram: 100, total_disk: 1000, available_disk: 600, model_count: 0, in_flight_count: 0 } + +const stableReport = { + row: { active: 'nats', epoch: 4, state: 'stable', changed_by: 'admin', changed_at: '2026-10-08T10:00:00Z', prev_epoch: 3 }, + active: 'nats', epoch: 4, state: 'stable', ok: true, + replicas: [ + { id: 'frontend-a', version: 'v4.12.0', ready_epoch: 4 }, + { id: 'frontend-b', version: 'v4.12.0', ready_epoch: 4 }, + ], + workers: [ + { id: 'w1', name: 'dual-worker', type: 'backend', attached: ['nats'], follow: ['nats', 'tunnel'], can_follow: true }, + { id: 'w2', name: 'tunnel-worker', type: 'backend', attached: ['tunnel'], follow: ['tunnel'], follow_error: 'no routable address: it cannot follow a change to NATS', can_follow: false }, + { id: 'w3', name: 'old-worker', type: 'backend', attached: [], can_follow: false }, + ], + in_flight: { loads: 0, jobs: 1, pending_claims: 2, claimed_claims: 0 }, + timings: {}, + settings: { nats_url: 'nats://nats:4222', nats_worker_url: '', prepare_timeout: '', transition_window: '', max_drain: '' }, +} + +const blockedDryRun = { + ...stableReport, target: 'tunnel', ok: false, + blockers: [ + { kind: 'worker', id: 'w3', reason: 'the worker predates carrier switching: it reports no capabilities and stays on NATS', forceable: true }, + ], +} +const cleanDryRun = { ...stableReport, target: 'tunnel', ok: true } + +async function mockCluster(page, { carrier = stableReport, carrierStatus = 200, nodes = [node] } = {}) { + await page.route('**/api/features', route => route.fulfill({ status: 200, contentType: 'application/json', body: JSON.stringify({ distributed: true }) })) + await page.route('**/api/auth/status', route => route.fulfill({ + status: 200, contentType: 'application/json', + body: JSON.stringify({ authEnabled: true, staticApiKeyRequired: false, providers: ['local'], user: { id: 'admin', name: 'Admin', role: 'admin', provider: 'local' } }), + })) + await page.route('**/api/nodes', route => route.fulfill({ status: 200, contentType: 'application/json', body: JSON.stringify(nodes) })) + const calls = { posts: [], puts: [] } + const state = { carrier } + await page.route('**/api/cluster/carrier', async route => { + const request = route.request() + if (request.method() === 'POST') { + const body = request.postDataJSON() + calls.posts.push(body) + const handler = state.onPost || (() => ({ status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5 } })) + const answer = handler(body) + return route.fulfill({ status: answer.status, contentType: 'application/json', body: JSON.stringify(answer.body) }) + } + return route.fulfill({ status: carrierStatus, contentType: 'application/json', body: JSON.stringify(carrierStatus === 200 ? state.carrier : { error: 'distributed mode is not enabled on this frontend' }) }) + }) + await page.route('**/api/cluster/settings', async route => { + const body = route.request().postDataJSON() + calls.puts.push(body) + const answer = (state.onPut || (() => ({ status: 200, body: { saved: true, nats: { reachable: true } } })))(body) + return route.fulfill({ status: answer.status, contentType: 'application/json', body: JSON.stringify(answer.body) }) + }) + return { calls, state } +} + +test.describe('Cluster transport panel', () => { + test('is hidden when the cluster routes answer 503', async ({ page }) => { + await mockCluster(page, { carrierStatus: 503 }) + await page.goto('/app/nodes') + await expect(page.getByLabel('Fleet health summary')).toBeVisible({ timeout: 15_000 }) + await expect(page.getByTestId('cluster-transport')).toHaveCount(0) + }) + + test('shows the active carrier, epoch, replicas and workers in the stable state', async ({ page }) => { + await mockCluster(page) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel).toBeVisible({ timeout: 15_000 }) + await expect(panel.getByTestId('transport-active')).toHaveText('NATS') + await expect(panel.getByTestId('transport-state')).toHaveText('stable') + await expect(panel.getByTestId('transport-epoch')).toHaveText('epoch 4') + await expect(panel.getByTestId('transport-replica')).toHaveCount(2) + await expect(panel.getByTestId('transport-replica').first()).toContainText('ready') + const workers = panel.getByTestId('transport-worker') + await expect(workers).toHaveCount(3) + await expect(workers.nth(0)).toContainText('Dual-capable') + await expect(workers.nth(1)).toContainText('Tunnel-only') + await expect(workers.nth(1)).toContainText('no routable address') + await expect(workers.nth(2)).toContainText('Legacy (NATS only)') + await expect(panel.getByRole('button', { name: /Abort change/ })).toHaveCount(0) + await expect(panel.getByLabel('NATS URL', { exact: true })).toHaveValue('nats://nats:4222') + }) + + test('dry run lists blockers and the confirmation gates the switch', async ({ page }) => { + const { calls, state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: blockedDryRun } : { status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5 } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + const start = dialog.getByRole('button', { name: /^Switch to Database tunnel$/ }) + await expect(start).toBeDisabled() + await dialog.getByRole('button', { name: 'Dry run' }).click() + await expect(dialog.getByTestId('transport-blocker')).toHaveCount(1) + await expect(dialog.getByTestId('transport-blocker')).toContainText('predates carrier switching') + await expect(dialog.getByTestId('transport-dry-verdict')).toContainText('1 blocker') + await expect(dialog.getByLabel('Live replicas').locator('li')).toHaveCount(2) + expect(calls.posts[0]).toEqual({ target: 'tunnel', dry_run: true }) + // A blocker without force keeps the button off, even once confirmed. + const confirm = dialog.getByLabel('I confirm these are all the frontends and all are upgraded') + await confirm.check() + await expect(start).toBeDisabled() + }) + + test('a clean dry run needs the confirmation, then starts the change', async ({ page }) => { + const { calls, state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: cleanDryRun } : { status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5 } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + await dialog.getByRole('button', { name: 'Dry run' }).click() + await expect(dialog.getByTestId('transport-dry-verdict')).toContainText('Nothing blocks') + const start = dialog.getByRole('button', { name: /^Switch to Database tunnel$/ }) + await expect(start).toBeDisabled() + await dialog.getByLabel('I confirm these are all the frontends and all are upgraded').check() + await expect(start).toBeEnabled() + await start.click() + await expect(dialog).toHaveCount(0) + expect(calls.posts.at(-1)).toEqual({ target: 'tunnel', force: false }) + }) + + test('force needs its own confirmation', async ({ page }) => { + const { calls, state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: blockedDryRun } : { status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5, force: true } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + await dialog.getByRole('button', { name: 'Dry run' }).click() + await dialog.getByLabel('I confirm these are all the frontends and all are upgraded').check() + await dialog.getByLabel('Force the change past the blockers above').check() + await dialog.getByRole('button', { name: /^Force switch to Database tunnel$/ }).click() + const confirm = page.getByRole('alertdialog') + await expect(confirm).toContainText('Force the change?') + // Nothing is sent until the second confirmation. + expect(calls.posts.filter(p => !p.dry_run)).toHaveLength(0) + await confirm.getByRole('button', { name: 'Force the change' }).click() + await expect(dialog).toHaveCount(0) + expect(calls.posts.at(-1)).toEqual({ target: 'tunnel', force: true }) + }) + + test('a refused request shows the blockers of the server', async ({ page }) => { + const { state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: cleanDryRun } : { status: 422, body: { error: 'the change is blocked', blockers: [{ kind: 'replica', id: 'frontend-b', reason: 'frontend-b cannot build the tunnel carrier', forceable: true }] } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + await dialog.getByRole('button', { name: 'Dry run' }).click() + await dialog.getByLabel('I confirm these are all the frontends and all are upgraded').check() + await dialog.getByRole('button', { name: /^Switch to Database tunnel$/ }).click() + await expect(dialog.getByRole('alert')).toContainText('frontend-b cannot build the tunnel carrier') + }) + + test('shows replica readiness and an abort button during prepare', async ({ page }) => { + const prepare = { + ...stableReport, state: 'prepare', epoch: 5, target: 'tunnel', + row: { ...stableReport.row, epoch: 5, state: 'prepare', target: 'tunnel' }, + replicas: [{ id: 'frontend-a', version: 'v4.12.0', ready_epoch: 5 }, { id: 'frontend-b', version: 'v4.12.0', ready_epoch: 4 }], + } + const { calls, state } = await mockCluster(page, { carrier: prepare }) + state.onPost = () => ({ status: 200, body: { active: 'nats', state: 'stable', epoch: 6 } }) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel.getByTestId('transport-state')).toHaveText('prepare') + await expect(panel.getByRole('button', { name: /^Switch to/ })).toBeDisabled() + const replicas = panel.getByTestId('transport-replica') + await expect(replicas.nth(0)).toContainText('ready') + await expect(replicas.nth(1)).toContainText('not ready') + await panel.getByRole('button', { name: 'Abort change' }).click() + await page.getByRole('alertdialog').getByRole('button', { name: 'Abort the change' }).click() + await expect.poll(() => calls.posts.length).toBe(1) + expect(calls.posts[0]).toEqual({ abort: true }) + }) + + test('counts down the drain after a commit', async ({ page }) => { + const draining = { + ...stableReport, active: 'tunnel', epoch: 6, + row: { ...stableReport.row, active: 'tunnel', epoch: 6, draining: 'nats' }, + drain_remaining_ns: 125e9, + } + await mockCluster(page, { carrier: draining }) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel.getByTestId('transport-active')).toHaveText('Database tunnel') + await expect(panel.getByTestId('transport-state')).toHaveText('draining') + const drain = panel.getByTestId('transport-drain') + await expect(drain).toContainText('Draining NATS') + await expect(drain).toContainText(/2:0\d left|2:1\d left|2:2\d left/) + const first = await drain.textContent() + await expect.poll(async () => drain.textContent(), { timeout: 5000 }).not.toBe(first) + }) + + test('saving the NATS settings reports reachability and does not switch', async ({ page }) => { + const { calls } = await mockCluster(page) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await panel.getByLabel('NATS URL', { exact: true }).fill('nats://nats.example:4222') + await panel.getByLabel('Worker NATS URL').fill('nats://nats.public.example:4222') + await panel.getByRole('button', { name: 'Save' }).click() + await expect(panel.getByTestId('transport-settings-result')).toContainText('reaches the NATS server') + await expect(panel.getByTestId('transport-settings-result')).toContainText('did not change') + expect(calls.puts[0]).toEqual({ nats_url: 'nats://nats.example:4222', nats_worker_url: 'nats://nats.public.example:4222' }) + expect(calls.posts).toHaveLength(0) + }) + + test('a NATS URL that cannot be reached shows the error', async ({ page }) => { + const { state } = await mockCluster(page) + state.onPut = () => ({ status: 422, body: { error: 'nats_url: this replica cannot reach the NATS server at nats://bad:4222: connection refused' } }) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await panel.getByLabel('NATS URL', { exact: true }).fill('nats://bad:4222') + await panel.getByRole('button', { name: 'Save' }).click() + await expect(panel.getByTestId('transport-settings-result')).toContainText('cannot reach the NATS server') + }) + + test('renders on a narrow viewport without horizontal overflow', async ({ page }) => { + await page.setViewportSize({ width: 390, height: 844 }) + await mockCluster(page) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel).toBeVisible({ timeout: 15_000 }) + const overflow = await page.evaluate(() => document.documentElement.scrollWidth - document.documentElement.clientWidth) + expect(overflow).toBeLessThanOrEqual(1) + }) +}) diff --git a/core/http/react-ui/src/App.css b/core/http/react-ui/src/App.css index 3ac4da254a34..4bad41b7d796 100644 --- a/core/http/react-ui/src/App.css +++ b/core/http/react-ui/src/App.css @@ -13875,3 +13875,35 @@ button.lane:hover { touch-action: none; cursor: grab; } + +/* Cluster transport panel (Nodes page) */ +.transport-panel { display: grid; gap: var(--spacing-md); } +.transport-head { align-items: flex-start; display: flex; flex-wrap: wrap; gap: var(--spacing-md); justify-content: space-between; } +.transport-head__line { align-items: center; display: flex; flex-wrap: wrap; gap: 10px; } +.transport-head__carrier { font-size: 1.375rem; letter-spacing: -.02em; } +.transport-head__actions { display: flex; flex-wrap: wrap; gap: var(--spacing-sm); } +.transport-grid { display: grid; gap: var(--spacing-lg); grid-template-columns: minmax(0, .8fr) minmax(0, 1.8fr); } +.transport-list { display: grid; gap: 6px; list-style: none; margin: 0 0 var(--spacing-md); padding: 0; } +.transport-list__row { align-items: center; display: flex; flex-wrap: wrap; font-size: var(--text-sm); gap: 8px; } +.transport-list__name { font-family: var(--font-mono); word-break: break-all; } +.transport-list__meta, .transport-list__state { color: var(--color-text-muted); font-size: var(--text-xs); } +.transport-dot { border-radius: 50%; display: inline-block; flex: none; height: 8px; width: 8px; } +.transport-dot--ok { background: var(--color-success); } +.transport-dot--warn { background: var(--color-warning); } +.transport-table-wrap { overflow-x: auto; } +.transport-table { border-collapse: collapse; font-size: var(--text-sm); width: 100%; } +.transport-table th { color: var(--color-text-muted); font-size: var(--text-xs); font-weight: 600; padding: 4px 10px 4px 0; text-align: left; } +.transport-table td { border-top: 1px solid var(--color-border-subtle); padding: 6px 10px 6px 0; vertical-align: top; } +.transport-table td:first-child { white-space: nowrap; } +.transport-error { color: var(--color-error); font-size: var(--text-sm); margin: 0; } +.transport-ok { color: var(--color-success); font-size: var(--text-sm); margin: 0; } +.transport-blockers { display: grid; gap: 8px; list-style: none; margin: 0 0 var(--spacing-md); padding: 0; } +.transport-blockers li { align-items: baseline; display: flex; flex-wrap: wrap; font-size: var(--text-sm); gap: 8px; } +.transport-settings { border-top: 1px solid var(--color-border-subtle); display: grid; gap: var(--spacing-sm); padding-top: var(--spacing-md); } +.transport-field { display: grid; font-size: var(--text-sm); gap: 4px; } +.transport-dialog { display: grid; gap: var(--spacing-md); padding: var(--spacing-lg); } +.transport-dialog__targets { justify-self: start; } +.transport-dialog__actions { display: flex; gap: var(--spacing-sm); } +.transport-dialog__result { display: grid; gap: var(--spacing-sm); } +.transport-dialog__footer { display: flex; flex-wrap: wrap; gap: var(--spacing-sm); justify-content: flex-end; } +@media (max-width: 760px) { .transport-grid { grid-template-columns: minmax(0, 1fr); } } diff --git a/core/http/react-ui/src/components/nodes/ClusterTransport.jsx b/core/http/react-ui/src/components/nodes/ClusterTransport.jsx new file mode 100644 index 000000000000..43f2d4e99984 --- /dev/null +++ b/core/http/react-ui/src/components/nodes/ClusterTransport.jsx @@ -0,0 +1,364 @@ +import { useCallback, useEffect, useRef, useState } from 'react' +import { clusterApi } from '../../utils/api' +import { + CARRIERS, PROFILE_LABELS, canForce, carrierLabel, formatCountdown, otherCarrier, + replicaReady, stateLabel, workerProfile, +} from '../../utils/clusterTransport' +import ConfirmDialog from '../ConfirmDialog' +import Modal from '../Modal' + +const POLL_IDLE_MS = 5000 +const POLL_BUSY_MS = 2000 + +// The routes answer 404 on a build without them and 503 on a frontend that is +// not distributed. Either way there is no transport to show. +function isUnavailable(error) { + return error?.status === 404 || error?.status === 503 || error?.status === 401 || error?.status === 403 +} + +function stateTone(state) { + if (state === 'stable') return 'badge-success' + if (state === 'draining') return 'badge-info' + return 'badge-warning' +} + +function ReplicaList({ replicas, report, label }) { + if (!replicas || replicas.length === 0) return

No live frontend replica reported.

+ return ( +
    + {replicas.map(replica => { + const ready = replicaReady(replica, report) + return ( +
  • +
    • + ) + })} +
+ ) +} + +function WorkerTable({ workers }) { + if (!workers || workers.length === 0) return

No worker is registered.

+ return ( +
+ + + + {workers.map(worker => ( + + + + + + + ))} + +
WorkerProfileAttached toFollow status
{worker.name || worker.id}{PROFILE_LABELS[workerProfile(worker)]}{(worker.attached || []).length ? worker.attached.map(carrierLabel).join(' + ') : 'not reported'}{worker.follow_error + ? {worker.follow_error} + : ok}
+
+ ) +} + +function BlockerList({ blockers }) { + return ( +
    + {blockers.map((blocker, index) => ( +
  • + {blocker.kind}{blocker.id ? ` ${blocker.id}` : ''} + {blocker.reason} + {blocker.forceable && can be forced} +
    • + ))} +
+ ) +} + +function SwitchDialog({ report, onClose, onStarted, addToast }) { + const [target, setTarget] = useState(otherCarrier(report.active)) + const [dry, setDry] = useState(null) + const [dryError, setDryError] = useState('') + const [running, setRunning] = useState(false) + const [confirmed, setConfirmed] = useState(false) + const [force, setForce] = useState(false) + const [confirmForce, setConfirmForce] = useState(false) + const [submitting, setSubmitting] = useState(false) + const [submitError, setSubmitError] = useState('') + + const resetResult = () => { setDry(null); setDryError(''); setConfirmed(false); setForce(false); setSubmitError('') } + const pickTarget = value => { setTarget(value); resetResult() } + + const runDry = async () => { + setRunning(true) + resetResult() + try { + setDry(await clusterApi.dryRun(target)) + } catch (error) { + setDryError(error.message || 'The dry run failed') + } finally { + setRunning(false) + } + } + + const blockers = dry?.blockers || [] + const blocked = blockers.length > 0 + const forceable = canForce(dry) + const sameCarrier = target === report.active + const mayStart = !!dry && confirmed && !sameCarrier && (!blocked || (force && forceable)) + + const start = async () => { + setSubmitting(true) + setSubmitError('') + try { + const answer = await clusterApi.switchCarrier(target, force) + addToast(`Change to ${carrierLabel(target)} started (epoch ${answer.epoch}).`, 'success') + onStarted() + onClose() + } catch (error) { + const refused = error?.body?.blockers + setSubmitError(Array.isArray(refused) && refused.length + ? `${error.message}: ${refused.map(b => b.reason).join('; ')}` + : (error.message || 'The request failed')) + } finally { + setSubmitting(false) + setConfirmForce(false) + } + } + + const onPrimary = () => { if (force) setConfirmForce(true); else start() } + const forcedNodes = blockers.filter(b => b.kind === 'worker').length + + return ( + <> + {} : onClose} maxWidth="680px" ariaLabel="Switch cluster transport"> +
+