From a3054442e2f9fb937d44648a52cf4725ea2625c8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 21:56:27 +0000 Subject: [PATCH 01/42] feat(cluster): record frontend replicas in a shared instances table Replicas need to find each other to relay worker traffic, and nothing in the tree recorded a replica's address. The advertised address is discovered by opening a UDP socket toward PostgreSQL and reading back the local address, which yields the interface every replica demonstrably shares without asking an operator to configure one. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/cluster_suite_test.go | 13 ++ core/services/cluster/instance.go | 189 ++++++++++++++++++++ core/services/cluster/instance_test.go | 101 +++++++++++ core/services/nodes/registry.go | 3 +- 4 files changed, 305 insertions(+), 1 deletion(-) create mode 100644 core/services/cluster/cluster_suite_test.go create mode 100644 core/services/cluster/instance.go create mode 100644 core/services/cluster/instance_test.go diff --git a/core/services/cluster/cluster_suite_test.go b/core/services/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..d821487298fc --- /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 Package Suite") +} diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go new file mode 100644 index 000000000000..7d4347204135 --- /dev/null +++ b/core/services/cluster/instance.go @@ -0,0 +1,189 @@ +// Package cluster records the frontend replicas that make up one LocalAI +// deployment and, later, the links between them. It is deliberately free of +// dependencies on core/services/nodes: nodes migrates and consumes the models +// declared here, so an import in the other direction would be a cycle. +package cluster + +import ( + "context" + "errors" + "fmt" + "net" + "net/url" + "strconv" + "strings" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// ErrInstanceNotFound reports that no row exists for the requested instance ID. +// Callers distinguish it from a transport failure to decide whether to +// re-register or to retry. +var ErrInstanceNotFound = errors.New("cluster: instance not found") + +// Instance is one live frontend replica, keyed by the ID that replica chose for +// itself. Column sizes mirror nodes.BackendNode so both tables agree on what an +// ID and a host:port look like. +type Instance struct { + ID string `gorm:"primaryKey;size:36" json:"id"` + AdvertisedAddr string `gorm:"size:255" json:"advertised_addr"` // host:port other replicas dial + Version string `gorm:"size:64" json:"version"` + LastSeen time.Time `gorm:"index" json:"last_seen"` +} + +// Registry reads and writes the instances table. +type Registry struct { + db *gorm.DB +} + +// NewRegistry returns a Registry over db. Migration is the caller's job; the +// nodes registry owns the AutoMigrate for every table in this deployment so +// that a single advisory lock covers them all. +func NewRegistry(db *gorm.DB) *Registry { + return &Registry{db: db} +} + +// Register records this replica's address, refreshing LastSeen. It upserts on +// the primary key rather than deleting and re-inserting, so a concurrent Live +// never observes a live replica as missing. +func (r *Registry) Register(ctx context.Context, id, addr, version string) error { + inst := Instance{ + ID: id, + AdvertisedAddr: addr, + Version: version, + LastSeen: time.Now(), + } + if err := r.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.AssignmentColumns([]string{"advertised_addr", "version", "last_seen"}), + }).Create(&inst).Error; err != nil { + return fmt.Errorf("registering instance %q: %w", id, err) + } + return nil +} + +// Heartbeat refreshes LastSeen for an already-registered instance. An unknown +// ID is an error rather than an insert: a heartbeat carries no address, so +// inserting would publish a replica nobody can reach. +func (r *Registry) Heartbeat(ctx context.Context, id string) error { + // gorm reports no error when a Where matches nothing, so the miss has to be + // read off RowsAffected. + res := r.db.WithContext(ctx).Model(&Instance{}). + Where("id = ?", id). + Update("last_seen", time.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 +} + +// Live returns the instances whose LastSeen is newer than now-within. +func (r *Registry) Live(ctx context.Context, within time.Duration) ([]Instance, error) { + var out []Instance + if err := r.db.WithContext(ctx). + Where("last_seen > ?", time.Now().Add(-within)). + Order("id"). + Find(&out).Error; err != nil { + return nil, fmt.Errorf("listing live instances: %w", err) + } + return out, nil +} + +// Get returns one instance, or ErrInstanceNotFound if it is not registered. +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 +} + +// 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. +// +// Failures are returned rather than papered over with a fallback such as +// 127.0.0.1, which would publish an address no peer can dial. +func DiscoverAdvertisedAddr(dsn string, port int) (string, error) { + 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 || local.IP.IsUnspecified() { + return "", fmt.Errorf("no local address on the route to database host %q", host) + } + return net.JoinHostPort(local.IP.String(), strconv.Itoa(port)), nil +} + +// 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/instance_test.go b/core/services/cluster/instance_test.go new file mode 100644 index 000000000000..112cc485123e --- /dev/null +++ b/core/services/cluster/instance_test.go @@ -0,0 +1,101 @@ +package cluster_test + +import ( + "context" + "net" + "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 + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) + reg = cluster.NewRegistry(db) + ctx = context.Background() + }) + + It("registers an instance and reads it back", func() { + Expect(reg.Register(ctx, "inst-a", "10.0.0.1:8080", "v1")).To(Succeed()) + + got, err := reg.Get(ctx, "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(got.AdvertisedAddr).To(Equal("10.0.0.1:8080")) + Expect(got.Version).To(Equal("v1")) + }) + + It("re-registering the same id updates the address instead of duplicating", func() { + Expect(reg.Register(ctx, "inst-a", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "inst-a", "10.0.0.9:9090", "v2")).To(Succeed()) + + live, err := reg.Live(ctx, time.Hour) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].AdvertisedAddr).To(Equal("10.0.0.9:9090")) + }) + + It("reports a missing instance distinguishably", func() { + _, err := reg.Get(ctx, "nope") + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + }) + + It("excludes instances whose heartbeat has aged out", func() { + Expect(reg.Register(ctx, "stale", "10.0.0.1:8080", "v1")).To(Succeed()) + // Age the row directly; sleeping in a spec is forbidden. + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "stale"). + Update("last_seen", time.Now().Add(-10*time.Minute)).Error).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(BeEmpty()) + }) + + It("brings a stale instance back with a heartbeat", func() { + Expect(reg.Register(ctx, "revive", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "revive"). + Update("last_seen", time.Now().Add(-10*time.Minute)).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, not a silent insert", func() { + Expect(reg.Heartbeat(ctx, "ghost")).To(MatchError(cluster.ErrInstanceNotFound)) + }) +}) + +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()) + }) +}) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 4b4f7f1c8b32..291ccd2dbb0f 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -9,6 +9,7 @@ import ( "github.com/google/uuid" "github.com/mudler/LocalAI/core/services/advisorylock" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/vrambudget" "github.com/mudler/xlog" @@ -442,7 +443,7 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}) + return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } From 6ef642f60d8059e0ee1df3cd3e044daf335e13f3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 22:08:42 +0000 Subject: [PATCH 02/42] fix(cluster): stamp liveness on the database clock and refuse undialable addresses DiscoverAdvertisedAddr promised to return an error rather than a fallback no peer can dial, but only rejected an unspecified address. With PostgreSQL on the same host or pod as a replica, which is compose, single-node and any sidecar layout, the route to it is loopback, so every replica advertised 127.0.0.1 and a peer dialling that reached itself. Loopback, link-local and zoned source addresses are now rejected with an error naming the remedy, and a port outside 1-65535 is rejected before it becomes an undialable address. Liveness was also measured on each replica's own clock: Register and Heartbeat stamped last_seen from the Go process, and Live compared those rows against the reading replica's time.Now(). Skew therefore shrank or stretched the window by writerBehind+readerAhead, evicting healthy peers or keeping dead ones. Both sides now use the database clock, which is the one clock every replica demonstrably shares. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/instance.go | 60 +++++++++++++++++++------- core/services/cluster/instance_test.go | 14 ++++++ 2 files changed, 59 insertions(+), 15 deletions(-) diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index 7d4347204135..027b0db3b0d9 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -49,16 +49,24 @@ func NewRegistry(db *gorm.DB) *Registry { // the primary key rather than deleting and re-inserting, so a concurrent Live // never observes a live replica as missing. func (r *Registry) Register(ctx context.Context, id, addr, version string) error { - inst := Instance{ - ID: id, - AdvertisedAddr: addr, - Version: version, - LastSeen: time.Now(), - } - if err := r.db.WithContext(ctx).Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "id"}}, - DoUpdates: clause.AssignmentColumns([]string{"advertised_addr", "version", "last_seen"}), - }).Create(&inst).Error; err != nil { + // last_seen is stamped by the database, never by this process. Liveness is + // compared across replicas, so it has to be measured on the one clock they + // all share; with per-replica clocks the effective Live window becomes + // `within - writerBehind - readerAhead`, which either evicts healthy peers + // or keeps dead ones alive. + if err := r.db.WithContext(ctx).Model(&Instance{}).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.Assignments(map[string]any{ + "advertised_addr": addr, + "version": version, + "last_seen": gorm.Expr("now()"), + }), + }).Create(map[string]any{ + "id": id, + "advertised_addr": addr, + "version": version, + "last_seen": gorm.Expr("now()"), + }).Error; err != nil { return fmt.Errorf("registering instance %q: %w", id, err) } return nil @@ -72,7 +80,7 @@ func (r *Registry) Heartbeat(ctx context.Context, id string) error { // read off RowsAffected. res := r.db.WithContext(ctx).Model(&Instance{}). Where("id = ?", id). - Update("last_seen", time.Now()) + Update("last_seen", gorm.Expr("now()")) if res.Error != nil { return fmt.Errorf("heartbeating instance %q: %w", id, res.Error) } @@ -85,8 +93,10 @@ func (r *Registry) Heartbeat(ctx context.Context, id string) error { // Live returns the instances whose LastSeen is newer than now-within. func (r *Registry) Live(ctx context.Context, within time.Duration) ([]Instance, error) { var out []Instance + // The cutoff is computed by the database for the same reason Register stamps + // there: a reader's clock must not decide whether another replica is alive. if err := r.db.WithContext(ctx). - Where("last_seen > ?", time.Now().Add(-within)). + Where("last_seen > now() - make_interval(secs => ?)", within.Seconds()). Order("id"). Find(&out).Error; err != nil { return nil, fmt.Errorf("listing live instances: %w", err) @@ -117,9 +127,19 @@ func (r *Registry) Get(ctx context.Context, id string) (*Instance, error) { // address to advertise. The caller supplies the port, since the frontend's // listening port has nothing to do with the database's. // -// Failures are returned rather than papered over with a fallback such as -// 127.0.0.1, which would publish an address no peer can dial. +// The discovery only holds while the database is a shared, remote host. When +// PostgreSQL runs on this same host or pod (compose, single-node, any sidecar +// layout) the route to it is loopback, and advertising 127.0.0.1 would make a +// peer dialling this replica reach itself instead. So an unspecified, loopback, +// or link-local 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 @@ -133,7 +153,17 @@ func DiscoverAdvertisedAddr(dsn string, port int) (string, error) { defer func() { _ = conn.Close() }() local, ok := conn.LocalAddr().(*net.UDPAddr) if !ok || local.IP == nil || local.IP.IsUnspecified() { - return "", fmt.Errorf("no local address on the route to database host %q", host) + return "", fmt.Errorf("no local address on the route to database host %q; set the advertised address explicitly", host) + } + if local.IP.IsLoopback() { + return "", fmt.Errorf("the route to database host %q is loopback (%s), so the database is local to this replica and its peer-reachable address cannot be discovered; set the advertised address explicitly", host, local.IP) + } + // A zone is only ever attached to a scoped (link-local) address, so this is + // the same rejection stated twice; the Zone check keeps the guarantee if a + // platform ever hands back a scoped address of another class, because + // IP.String() would silently drop the %iface and yield an undialable host. + if local.IP.IsLinkLocalUnicast() || local.Zone != "" { + return "", fmt.Errorf("the route to database host %q is link-local (%s), which peers on other hosts cannot dial; set the advertised address explicitly", host, local.IP) } return net.JoinHostPort(local.IP.String(), strconv.Itoa(port)), nil } diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go index 112cc485123e..8e74e930e3b9 100644 --- a/core/services/cluster/instance_test.go +++ b/core/services/cluster/instance_test.go @@ -98,4 +98,18 @@ var _ = Describe("Advertised address discovery", func() { _, err := cluster.DiscoverAdvertisedAddr("", 8080) Expect(err).To(HaveOccurred()) }) + + // A database on this same host routes over loopback on every platform, so + // this is deterministic rather than host-dependent. 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"))) + }) }) From 6ce3edaa43a82d66309e629e39d8c41680cc78c0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 22:19:05 +0000 Subject: [PATCH 03/42] feat(cluster): add the bidirectional splice used by the relay and tunnel Returns on the first direction to finish and closes both sides so the other unblocks; a sequential copy deadlocks on any protocol where the far side speaks first. EOF and use-of-closed are normal termination, not errors. The fourth spec covers a peer that stops reading mid-body, the case where a copy is parked in Write rather than in Read. The other three tear down an idle splice and pass even against a Splice that closes only one side. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/splice.go | 68 ++++++++++++++++ core/services/cluster/splice_test.go | 114 +++++++++++++++++++++++++++ 2 files changed, 182 insertions(+) create mode 100644 core/services/cluster/splice.go create mode 100644 core/services/cluster/splice_test.go diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go new file mode 100644 index 000000000000..a31f0475109c --- /dev/null +++ b/core/services/cluster/splice.go @@ -0,0 +1,68 @@ +package cluster + +import ( + "errors" + "io" + "net" +) + +// Splice joins two streams and copies bytes between them in both directions +// until one direction finishes, then closes both so the other unblocks and +// returns. It is the primitive under the inter-replica relay and the worker +// tunnel, so it carries gRPC: both directions can be live at once and either +// peer may speak first, which is why the copies run concurrently. A sequential +// io.Copy then io.Copy would deadlock waiting for a request on a stream whose +// far side is waiting for a response. +// +// The error reported is the one from the direction that finished first, with +// the endings that mean "someone closed" mapped to nil. The other direction's +// error is discarded: by then it is only an echo of the Close done here. +func Splice(a, b io.ReadWriteCloser) error { + errs := make(chan error, 2) + go func() { errs <- copyStream(b, a) }() + go func() { errs <- copyStream(a, b) }() + + first := <-errs + + // Closing both ends is what releases the other direction, whether it is + // parked in Read or halfway through a Write nobody is draining. Each end + // is closed exactly once, here and nowhere else, so a stream that reports + // an error on a second Close (yamux does) never sees one. + closeErrA := a.Close() + closeErrB := b.Close() + + // Wait for the second direction so no copy is still touching either + // stream once Splice has returned. + <-errs + + if first != nil { + return first + } + // A close that fails on a stream that was otherwise healthy is worth + // reporting; a close of an already-dead stream is not. + if err := normalizeStreamErr(closeErrA); err != nil { + return err + } + return normalizeStreamErr(closeErrB) +} + +// copyStream moves one direction and reports only genuine transport failures. +func copyStream(dst io.Writer, src io.Reader) error { + _, err := io.Copy(dst, src) + return normalizeStreamErr(err) +} + +// normalizeStreamErr drops the endings that mean the conversation is over +// rather than broken. io.EOF is the clean end of a stream, net.ErrClosed is +// what a socket or yamux stream reports once it or its peer has been closed, +// and io.ErrClosedPipe is the same condition on an in-memory pipe. Anything +// else is a real failure the caller should see. +func normalizeStreamErr(err error) error { + if err == nil || + errors.Is(err, io.EOF) || + errors.Is(err, net.ErrClosed) || + errors.Is(err, io.ErrClosedPipe) { + return nil + } + return err +} diff --git a/core/services/cluster/splice_test.go b/core/services/cluster/splice_test.go new file mode 100644 index 000000000000..9e3529e65e1a --- /dev/null +++ b/core/services/cluster/splice_test.go @@ -0,0 +1,114 @@ +package cluster_test + +import ( + "errors" + "io" + "net" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Splice", func() { + // pipePair returns two connected in-memory conns. + newPair := func() (net.Conn, net.Conn) { return net.Pipe() } + + It("copies bytes in both directions", func() { + aLeft, aRight := newPair() + bLeft, bRight := newPair() + + done := make(chan error, 1) + go func() { done <- cluster.Splice(aRight, bLeft) }() + + go func() { + _, _ = aLeft.Write([]byte("ping")) + }() + buf := make([]byte, 4) + Expect(bRight.SetReadDeadline(time.Now().Add(5 * time.Second))).To(Succeed()) + _, err := io.ReadFull(bRight, buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("ping")) + + go func() { + _, _ = bRight.Write([]byte("pong")) + }() + Expect(aLeft.SetReadDeadline(time.Now().Add(5 * time.Second))).To(Succeed()) + _, err = io.ReadFull(aLeft, buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("pong")) + + Expect(aLeft.Close()).To(Succeed()) + Eventually(done, "5s").Should(Receive()) + }) + + It("returns when one side closes, and closes the other", func() { + aLeft, aRight := newPair() + bLeft, bRight := newPair() + + done := make(chan error, 1) + go func() { done <- cluster.Splice(aRight, bLeft) }() + + Expect(aLeft.Close()).To(Succeed()) + Eventually(done, "5s").Should(Receive(BeNil())) + + // The far side must have been closed too, so a read there fails + // rather than blocking forever. The read runs in a goroutine and is + // polled instead of carrying a read deadline: net.Pipe refuses to set + // a deadline once *either* end is closed, so a deadline here would + // fail exactly when Splice did its job. + reads := make(chan error, 1) + go func() { + _, err := bRight.Read(make([]byte, 1)) + reads <- err + }() + var err error + Eventually(reads, "5s").Should(Receive(&err)) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, io.EOF)).To(BeTrue()) + }) + + It("does not leak a goroutine when both sides close", func() { + aLeft, aRight := newPair() + bLeft, bRight := newPair() + + done := make(chan error, 1) + go func() { done <- cluster.Splice(aRight, bLeft) }() + + Expect(aLeft.Close()).To(Succeed()) + Expect(bRight.Close()).To(Succeed()) + Eventually(done, "5s").Should(Receive()) + }) + + // The three specs above only ever tear down an idle splice: at the moment + // of Close no copy is parked inside a Write. A relayed inference response + // is the opposite case, a reader that walks away mid-body while 50MB is + // still being pushed at it, so this covers the direction that is blocked + // in Write rather than in Read when its peer disappears. + It("returns when the reader disappears while a write is in flight", func() { + aLeft, aRight := newPair() + bLeft, bRight := newPair() + + done := make(chan error, 1) + go func() { done <- cluster.Splice(aRight, bLeft) }() + + // Nothing ever reads from bRight, so the a->b direction parks inside + // Write on an unbuffered pipe with the payload half-delivered. + payload := make([]byte, 1<<20) + writes := make(chan error, 1) + go func() { + _, err := aLeft.Write(payload) + writes <- err + }() + + Expect(bRight.Close()).To(Succeed()) + Eventually(done, "5s").Should(Receive()) + + // The abandoned writer must be released as well, and only Splice + // closing its end can do that: no deadline is set on aLeft, so a + // splice that forgot to close would leave this write parked forever. + Eventually(writes, "5s").Should(Receive(HaveOccurred())) + }) +}) From 77e204124944b8a512eac0345875b839605ce80d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 22:33:54 +0000 Subject: [PATCH 04/42] fix(cluster): treat a yamux teardown in Splice as a normal ending go-yamux/v5 matches none of its errors against net.ErrClosed, so the classifier reported an ordinary teardown as a failure: when the session has gone away, the FIN that Splice's own Close writes returns ErrSessionShutdown, and a stream torn down under a live copy surfaces as ErrStreamClosed or a reset. Splice owns that Close, so it owns the errors it produces; the sentinels are named here rather than injected by the caller, which would make a forgotten classifier reintroduce the same bug silently. Cover the error half of the contract, which no in-memory pipe could reach: a scripted stream now feeds Splice a genuine transport failure and each closed-stream ending in turn. Replacing the tail of Splice with "return nil" passed every previous spec. Also assert that Splice does not return until the second direction has finished, rename a spec that promised a leak check it never made, and correct two comments that claimed more than the code did. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/splice.go | 43 ++++++-- core/services/cluster/splice_test.go | 145 ++++++++++++++++++++++++++- go.mod | 2 +- 3 files changed, 181 insertions(+), 9 deletions(-) diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go index a31f0475109c..3f7bd8fbf7d1 100644 --- a/core/services/cluster/splice.go +++ b/core/services/cluster/splice.go @@ -4,6 +4,8 @@ import ( "errors" "io" "net" + + "github.com/libp2p/go-yamux/v5" ) // Splice joins two streams and copies bytes between them in both directions @@ -14,9 +16,16 @@ import ( // io.Copy then io.Copy would deadlock waiting for a request on a stream whose // far side is waiting for a response. // +// EOF in one direction therefore truncates whatever is still in flight in the +// other. That is right for gRPC, HTTP/2 and yamux, which end a stream in both +// directions at once, but a future caller relaying raw TCP with a half-close +// would lose the response body still arriving after the request's CloseWrite. +// // The error reported is the one from the direction that finished first, with // the endings that mean "someone closed" mapped to nil. The other direction's -// error is discarded: by then it is only an echo of the Close done here. +// error is dropped; most of the time it is an echo of the Close below, but it +// can also be a genuine failure that lost the race, so a Splice error means +// "one direction failed", never "only this failed". func Splice(a, b io.ReadWriteCloser) error { errs := make(chan error, 2) go func() { errs <- copyStream(b, a) }() @@ -31,8 +40,13 @@ func Splice(a, b io.ReadWriteCloser) error { closeErrA := a.Close() closeErrB := b.Close() - // Wait for the second direction so no copy is still touching either - // stream once Splice has returned. + // Wait for the second direction so no copy is still touching either stream + // once Splice has returned. This is load-bearing: it assumes Close unblocks + // a copy parked in Read or Write, and a stream where that is false hangs + // here rather than leaking a goroutine. Both callers satisfy it. net.Conn + // does, and so does go-yamux/v5, whose Close sets readErr and calls + // notifyWaiting to wake a parked Read while a parked Write returns + // ErrStreamClosed. <-errs if first != nil { @@ -54,14 +68,29 @@ func copyStream(dst io.Writer, src io.Reader) error { // normalizeStreamErr drops the endings that mean the conversation is over // rather than broken. io.EOF is the clean end of a stream, net.ErrClosed is -// what a socket or yamux stream reports once it or its peer has been closed, -// and io.ErrClosedPipe is the same condition on an in-memory pipe. Anything -// else is a real failure the caller should see. +// what a socket reports once it or its peer has been closed, and +// io.ErrClosedPipe is the same condition on an in-memory pipe. +// +// The yamux sentinels are here because Splice owns the Close that produces +// them: closing a stream whose session has already gone away returns +// ErrSessionShutdown from the FIN write, and a stream torn down under a live +// copy surfaces as ErrStreamClosed or a reset. None of yamux's error types +// match net.ErrClosed, so each has to be named. Matching ErrStreamReset covers +// every *StreamError and *GoAwayError, both of which report themselves as a +// reset; ErrStreamClosed is a plain sentinel and matches only itself. +// +// Anything else is a real failure the caller should see. Note that a peer +// resetting a socket mid-stream (ECONNRESET, EPIPE) is deliberately not in +// this set: whether an abandoned request is routine is the caller's policy, +// not this primitive's. func normalizeStreamErr(err error) error { if err == nil || errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || - errors.Is(err, io.ErrClosedPipe) { + errors.Is(err, io.ErrClosedPipe) || + errors.Is(err, yamux.ErrStreamClosed) || + errors.Is(err, yamux.ErrStreamReset) || + errors.Is(err, yamux.ErrSessionShutdown) { return nil } return err diff --git a/core/services/cluster/splice_test.go b/core/services/cluster/splice_test.go index 9e3529e65e1a..5b69eef720ef 100644 --- a/core/services/cluster/splice_test.go +++ b/core/services/cluster/splice_test.go @@ -4,8 +4,12 @@ import ( "errors" "io" "net" + "sync" + "sync/atomic" "time" + "github.com/libp2p/go-yamux/v5" + "github.com/mudler/LocalAI/core/services/cluster" . "github.com/onsi/ginkgo/v2" @@ -70,7 +74,7 @@ var _ = Describe("Splice", func() { Expect(errors.Is(err, io.EOF)).To(BeTrue()) }) - It("does not leak a goroutine when both sides close", func() { + It("returns when both sides close", func() { aLeft, aRight := newPair() bLeft, bRight := newPair() @@ -111,4 +115,143 @@ var _ = Describe("Splice", func() { // splice that forgot to close would leave this write parked forever. Eventually(writes, "5s").Should(Receive(HaveOccurred())) }) + // net.Pipe can only ever end in EOF or a closed pipe, so the error half of + // the contract needs a stream that can be told how to fail. + Context("when a stream fails rather than closing", func() { + errBoom := errors.New("transport exploded") + + It("reports a genuine transport error", func() { + failing := &scriptedStream{readErr: errBoom} + idle := &scriptedStream{} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(failing, idle) }() + + var err error + Eventually(done, "5s").Should(Receive(&err)) + Expect(errors.Is(err, errBoom)).To(BeTrue()) + + // Closed exactly once each: a second Close is what makes a yamux + // stream complain about a teardown that went fine. + Expect(failing.closes()).To(Equal(int32(1))) + Expect(idle.closes()).To(Equal(int32(1))) + }) + + It("reports a genuine failure from its own Close", func() { + failing := &scriptedStream{closeErr: errBoom} + idle := &scriptedStream{} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(idle, failing) }() + + Expect(idle.Close()).To(Succeed()) + var err error + Eventually(done, "5s").Should(Receive(&err)) + Expect(errors.Is(err, errBoom)).To(BeTrue()) + }) + + // Every one of these means "a stream we were copying through was + // closed". The yamux entries are the teardown Splice itself provokes: + // none of them matches net.ErrClosed, so each has to be classified by + // name or a normal relayed request ends up reported as a failure. + DescribeTable("treats a closed stream as normal termination", + func(ending error) { + ended := &scriptedStream{readErr: ending} + idle := &scriptedStream{} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(ended, idle) }() + + Eventually(done, "5s").Should(Receive(BeNil())) + }, + Entry("EOF", io.EOF), + Entry("a closed socket", net.ErrClosed), + Entry("a closed in-memory pipe", io.ErrClosedPipe), + Entry("a closed yamux stream", yamux.ErrStreamClosed), + Entry("a reset yamux stream", yamux.ErrStreamReset), + Entry("a shut-down yamux session", yamux.ErrSessionShutdown), + Entry("a stream reset by the remote", &yamux.StreamError{ErrorCode: 1, Remote: true}), + Entry("a go-away from the remote", yamux.ErrRemoteGoAway), + ) + + // The tunnel's own teardown: the local backend finishes normally while + // the yamux session has already gone away, so the FIN that Splice's + // Close writes fails. Nothing went wrong and nothing may be reported. + It("does not report a shut-down session on its own Close", func() { + stream := &scriptedStream{closeErr: yamux.ErrSessionShutdown} + backend := &scriptedStream{readErr: io.EOF} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(stream, backend) }() + + Eventually(done, "5s").Should(Receive(BeNil())) + }) + + // The anti-leak guarantee, which the pipe specs cannot see because + // their parked copy is released too quickly to catch Splice in the + // act. Waking a copy is asynchronous on a real stream (yamux's Close + // notifies the reader, which then has to be scheduled), so this stream + // splits the two: Close records itself, and the spec decides when the + // parked Read actually returns. + It("does not return until the second direction has finished", func() { + parked := &scriptedStream{holdReadPastClose: true} + ending := &scriptedStream{readErr: io.EOF} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(ending, parked) }() + + Eventually(parked.closes, "5s").Should(Equal(int32(1))) + Consistently(done, "200ms").ShouldNot(Receive()) + + parked.release() + Eventually(done, "5s").Should(Receive(BeNil())) + }) + }) }) + +// scriptedStream is an io.ReadWriteCloser whose endings the spec dictates, so +// Splice can be fed failures no in-memory pipe can produce. With no readErr it +// parks in Read until Close, standing in for an idle half of a live stream. +type scriptedStream struct { + readErr error + closeErr error + // holdReadPastClose keeps a parked Read blocked until release is called, + // standing in for the gap between a Close waking a reader and that reader + // running. Without it, Close releases the Read as a real stream does. + holdReadPastClose bool + + releaseOnce sync.Once + released chan struct{} + initOnce sync.Once + closeN atomic.Int32 +} + +func (s *scriptedStream) gate() chan struct{} { + s.initOnce.Do(func() { s.released = make(chan struct{}) }) + return s.released +} + +func (s *scriptedStream) release() { + gate := s.gate() + s.releaseOnce.Do(func() { close(gate) }) +} + +func (s *scriptedStream) Read(p []byte) (int, error) { + if s.readErr != nil { + return 0, s.readErr + } + <-s.gate() + return 0, io.EOF +} + +func (s *scriptedStream) Write(p []byte) (int, error) { return len(p), nil } + +func (s *scriptedStream) Close() error { + s.closeN.Add(1) + if !s.holdReadPastClose { + s.release() + } + return s.closeErr +} + +func (s *scriptedStream) closes() int32 { return s.closeN.Load() } diff --git a/go.mod b/go.mod index 5ed7e0b515d7..c3ded1cc3712 100644 --- a/go.mod +++ b/go.mod @@ -32,6 +32,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 @@ -321,7 +322,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 8ebc24194c6111f544acedc3335f490b9d90d54b Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 22:49:57 +0000 Subject: [PATCH 05/42] fix(cluster): report a dead yamux session instead of swallowing it Matching yamux errors with errors.Is was too broad. Session.close hands every live stream ErrStreamReset wrapped around whatever killed the connection, so a keepalive timeout, a broken TCP connection or a peer that simply vanished all matched, and a relayed request that died reported a clean ending. Nothing upstream would have retried or logged it. Match the plain sentinels by identity, since only identity separates a stream that was reset from the wrapped form that means the session died. Treat a StreamError as a per-stream reset, and a GoAwayError as normal only when it carries the no-error code, read off ErrRemoteGoAway because the constant is unexported. ErrSessionShutdown needs no entry of its own; it is a GoAwayError with that code. Order matters as much as the matching: session death wraps its cause, which is routinely io.EOF or a closed socket, so the mux checks run before the generic endings. Reversing them alone puts a vanished peer back to nil. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/splice.go | 93 ++++++++++++++++++++++------ core/services/cluster/splice_test.go | 67 ++++++++++++++++++++ 2 files changed, 140 insertions(+), 20 deletions(-) diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go index 3f7bd8fbf7d1..ba5ac28af294 100644 --- a/core/services/cluster/splice.go +++ b/core/services/cluster/splice.go @@ -35,8 +35,10 @@ func Splice(a, b io.ReadWriteCloser) error { // Closing both ends is what releases the other direction, whether it is // parked in Read or halfway through a Write nobody is draining. Each end - // is closed exactly once, here and nowhere else, so a stream that reports - // an error on a second Close (yamux does) never sees one. + // is closed exactly once, here and nowhere else, which keeps the error + // below meaningful: a second Close of a yamux stream that was reset + // returns the error that killed it, and Splice would have no way to tell + // that from a fresh failure. closeErrA := a.Close() closeErrB := b.Close() @@ -71,27 +73,78 @@ func copyStream(dst io.Writer, src io.Reader) error { // what a socket reports once it or its peer has been closed, and // io.ErrClosedPipe is the same condition on an in-memory pipe. // -// The yamux sentinels are here because Splice owns the Close that produces -// them: closing a stream whose session has already gone away returns -// ErrSessionShutdown from the FIN write, and a stream torn down under a live -// copy surfaces as ErrStreamClosed or a reset. None of yamux's error types -// match net.ErrClosed, so each has to be named. Matching ErrStreamReset covers -// every *StreamError and *GoAwayError, both of which report themselves as a -// reset; ErrStreamClosed is a plain sentinel and matches only itself. -// -// Anything else is a real failure the caller should see. Note that a peer -// resetting a socket mid-stream (ECONNRESET, EPIPE) is deliberately not in -// this set: whether an abandoned request is routine is the caller's policy, -// not this primitive's. +// The mux checks run first, and that ordering is load-bearing: a dying yamux +// session hands every live stream its own cause wrapped up (session.go:330), +// and that cause is routinely io.EOF or a closed-socket error, so consulting +// the generic endings first would report a peer that vanished mid-request as a +// clean completion. func normalizeStreamErr(err error) error { - if err == nil || - errors.Is(err, io.EOF) || + if err == nil { + return nil + } + if isMuxSessionFailure(err) { + return err + } + if isMuxStreamTeardown(err) { + return nil + } + if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || - errors.Is(err, io.ErrClosedPipe) || - errors.Is(err, yamux.ErrStreamClosed) || - errors.Is(err, yamux.ErrStreamReset) || - errors.Is(err, yamux.ErrSessionShutdown) { + errors.Is(err, io.ErrClosedPipe) { return nil } return err } + +// normalGoAwayCode is yamux's "no error" go-away code, read off a sentinel +// declared with it because the constant itself is unexported. +var normalGoAwayCode = yamux.ErrRemoteGoAway.ErrorCode + +// isMuxSessionFailure reports whether err is the yamux session underneath a +// stream dying, as opposed to a single stream being torn down. The distinction +// matters because Splice must stay quiet about the teardown it provokes itself +// while still reporting a dead peer: a keepalive timeout, a broken TCP +// connection or a protocol error under a relayed request has to reach the +// caller, or a failed inference looks like a finished one. +func isMuxSessionFailure(err error) bool { + // A go-away ends the whole session. Only the "no error" code is a normal + // ending; a protocol or internal error go-away is a real failure. + var goAway *yamux.GoAwayError + if errors.As(err, &goAway) { + return goAway.ErrorCode != normalGoAwayCode + } + // A stream error is scoped to one stream, whatever killed it. + var streamErr *yamux.StreamError + if errors.As(err, &streamErr) { + return false + } + // Session.close gives every stream it kills ErrStreamReset wrapped around + // the cause, so the bare sentinel means this stream was reset and a + // wrapped one means the session died under it. Identity is what separates + // them; errors.Is cannot. + return errors.Is(err, yamux.ErrStreamReset) && err != yamux.ErrStreamReset +} + +// isMuxStreamTeardown reports whether err is yamux ending one stream, which +// Splice mostly provokes itself: closing a stream whose session has already +// shut down normally returns ErrSessionShutdown from the FIN write, and a copy +// parked on a stream that gets closed comes back with ErrStreamClosed or a +// reset. A reset does not only mean that, though; the same sentinel heads the +// error a dying session hands its streams, which is why isMuxSessionFailure +// runs first. None of yamux's error types match net.ErrClosed, so all of this +// has to be recognised by shape. +func isMuxStreamTeardown(err error) bool { + // Sentinels by identity, never errors.Is: the wrapped forms belong to a + // dead session and are reported instead. ErrSessionShutdown is absent on + // purpose rather than by oversight, being a *GoAwayError carrying the + // normal code, which the last check below covers. + if err == yamux.ErrStreamClosed || err == yamux.ErrStreamReset { + return true + } + var streamErr *yamux.StreamError + if errors.As(err, &streamErr) { + return true + } + var goAway *yamux.GoAwayError + return errors.As(err, &goAway) && goAway.ErrorCode == normalGoAwayCode +} diff --git a/core/services/cluster/splice_test.go b/core/services/cluster/splice_test.go index 5b69eef720ef..f18618b6b2cd 100644 --- a/core/services/cluster/splice_test.go +++ b/core/services/cluster/splice_test.go @@ -2,6 +2,7 @@ package cluster_test import ( "errors" + "fmt" "io" "net" "sync" @@ -174,6 +175,72 @@ var _ = Describe("Splice", func() { Entry("a go-away from the remote", yamux.ErrRemoteGoAway), ) + // sessionDeath is the exact shape Session.close hands every live + // stream when the session dies for a non-go-away reason + // (session.go:330). It matters that these are wrapped: the cause it + // carries is routinely io.EOF or a closed socket, so a classifier that + // looked at the cause would call a vanished peer a clean ending. + sessionDeath := func(cause error) error { + return fmt.Errorf("%w: connection closed: %w", yamux.ErrStreamReset, cause) + } + + // A dead peer under a relayed inference request has to reach the + // caller. If it arrives as nil, a failed request looks like a finished + // one and nothing upstream retries or logs it. + DescribeTable("reports the session dying under a stream", + func(ending error) { + dead := &scriptedStream{readErr: ending} + idle := &scriptedStream{} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(dead, idle) }() + + var err error + Eventually(done, "5s").Should(Receive(&err)) + Expect(err).To(MatchError(ending)) + }, + Entry("a keepalive timeout", sessionDeath(yamux.ErrKeepAliveTimeout)), + Entry("a broken connection", sessionDeath(errors.New("read tcp 10.0.0.1:4000: broken pipe"))), + Entry("a peer that vanished", sessionDeath(io.EOF)), + Entry("a protocol-error go-away", &yamux.GoAwayError{Remote: true, ErrorCode: 1}), + Entry("an internal-error go-away", &yamux.GoAwayError{Remote: true, ErrorCode: 2}), + ) + + // The distinction the classifier turns on, in one spec: yamux uses the + // same sentinel for "this stream was reset", which Splice provokes + // itself and must stay quiet about, and as the head of the wrapped + // error meaning "the session died", which it must report. Only + // identity separates them. + It("separates a bare reset from a session that died wrapping one", func() { + spliceEnding := func(ending error) error { + done := make(chan error, 1) + go func() { + done <- cluster.Splice(&scriptedStream{readErr: ending}, &scriptedStream{}) + }() + var err error + EventuallyWithOffset(1, done, "5s").Should(Receive(&err)) + return err + } + + Expect(spliceEnding(yamux.ErrStreamReset)).To(BeNil()) + Expect(spliceEnding(sessionDeath(yamux.ErrKeepAliveTimeout))).ToNot(BeNil()) + }) + + // Session death also arrives through the Close Splice makes itself, on + // a stream whose session died while the other side was finishing. That + // is not the quiet teardown ErrSessionShutdown describes. + It("reports a session that died, even from its own Close", func() { + stream := &scriptedStream{closeErr: sessionDeath(yamux.ErrKeepAliveTimeout)} + backend := &scriptedStream{readErr: io.EOF} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(stream, backend) }() + + var err error + Eventually(done, "5s").Should(Receive(&err)) + Expect(err).To(MatchError(yamux.ErrKeepAliveTimeout)) + }) + // The tunnel's own teardown: the local backend finishes normally while // the yamux session has already gone away, so the FIN that Splice's // Close writes fails. Nothing went wrong and nothing may be reported. From e957ff1ca255ebc327599dafba7e943a8df34785 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 23:14:37 +0000 Subject: [PATCH 06/42] feat(cluster): accept authenticated peer links on /api/cluster/peer Upgrades to a WebSocket, wraps it as a yamux server session and hands it to the caller. Rejects before upgrading so an unauthenticated dial sees a 401 rather than a WebSocket error, which is what the route-coverage test asserts. The adapter keeps the reader of a partially consumed message across Read calls. yamux reads through a 4 KiB bufio.Reader, so a small-payload test cannot see a dropped message tail; the framing specs drive the adapter directly with buffers smaller than the message. An empty configured token authorizes nobody here, unlike the worker file transfer server's check: this route is registered in every deployment, so failing open would publish an unauthenticated mux. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/http/auth/public_routes.go | 9 +- .../endpoints/cluster/cluster_suite_test.go | 13 ++ core/http/endpoints/cluster/peer.go | 104 +++++++++ core/http/endpoints/cluster/peer_test.go | 117 ++++++++++ core/http/endpoints/cluster/wsconn.go | 129 +++++++++++ core/http/endpoints/cluster/wsconn_test.go | 216 ++++++++++++++++++ 6 files changed, 587 insertions(+), 1 deletion(-) create mode 100644 core/http/endpoints/cluster/cluster_suite_test.go create mode 100644 core/http/endpoints/cluster/peer.go create mode 100644 core/http/endpoints/cluster/peer_test.go create mode 100644 core/http/endpoints/cluster/wsconn.go create mode 100644 core/http/endpoints/cluster/wsconn_test.go diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index 658205a78f8f..e1db88a02f60 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -5,6 +5,8 @@ package auth import ( "net/http" "strings" + + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" ) type publicRouteRule struct { @@ -77,5 +79,10 @@ func isPublicRoute(method, path string) bool { // 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/") + // The peer link carries the cluster token in an Authorization header that + // no browser session ever sets, and its handler checks that token itself. + // The prefix comes from the endpoints package so the route and the + // exemption cannot drift apart. + return strings.HasPrefix(path, "/api/node/") || + strings.HasPrefix(path, clusterep.AlternativeAuthPrefix) } 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..00f45e37a672 --- /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 TestClusterEndpoints(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster Endpoints Suite") +} diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go new file mode 100644 index 000000000000..4b6ef542caf4 --- /dev/null +++ b/core/http/endpoints/cluster/peer.go @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: MIT + +// Package cluster serves the replica-to-replica link that a LocalAI frontend +// uses to reach a worker tunnel it does not own. A peer dials +// GET /api/cluster/peer, the connection becomes one multiplexed yamux session, +// and the relay opens a stream on it per request. +package cluster + +import ( + "crypto/subtle" + "net/http" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + "github.com/mudler/xlog" +) + +// AlternativeAuthPrefix is the path prefix whose credentials are checked by +// this package rather than by the global session middleware. The auth layer +// consults this same constant, so the two cannot drift apart and leave every +// peer dial answering 401. +const AlternativeAuthPrefix = "/api/cluster/" + +// PeerPath is the route a peer replica dials. +const PeerPath = AlternativeAuthPrefix + "peer" + +// RegisterClusterRoutes registers the peer link. onPeer receives every +// authenticated session; see PeerHandler for what it is expected to do with it. +func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { + e.GET(PeerPath, PeerHandler(token, onPeer)) +} + +// PeerHandler upgrades an authenticated peer dial to a WebSocket, wraps it as +// a yamux server session and hands it to onSession. +// +// onSession runs on the request goroutine, so it must return promptly; the +// session outlives the handler because the upgrade hijacks the connection, and +// closing it is the caller's job. +func PeerHandler(token string, onSession func(peerID string, sess *yamux.Session)) echo.HandlerFunc { + // gorilla's default CheckOrigin already restricts a browser to same-origin + // and lets a header-less client (which every peer is) through, so the + // zero value is what this link wants. + upgrader := websocket.Upgrader{} + + return func(c echo.Context) error { + // Reject before upgrading. Upgrading and then closing would give the + // dialer a WebSocket error in place of an HTTP status, and both the + // route-coverage test and a peer's own retry logic read the status. + if !authorizedPeer(c.Request(), token) { + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + peerID := c.QueryParam("id") + if peerID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing peer id") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has already written its own failure to the client. + xlog.Debug("cluster peer link upgrade failed", "peer", peerID, "error", err) + return nil + } + + // Server side of the mux: the dialing peer is the client, so it owns + // the odd stream IDs and this side the even ones. + sess, err := yamux.Server(WebsocketConn(ws), nil, nil) + if err != nil { + xlog.Error("cluster peer link session setup failed", "peer", peerID, "error", err) + _ = ws.Close() + return nil + } + + if onSession == nil { + // Nothing will ever read from this session, so do not leave the + // peer believing it has a live link. + _ = sess.Close() + return nil + } + + xlog.Debug("cluster peer link established", "peer", peerID, "remote", ws.RemoteAddr().String()) + onSession(peerID, sess) + return nil + } +} + +// authorizedPeer compares the request's bearer token with the cluster token in +// constant time, matching the check the worker file-transfer server makes. +// +// Unlike that one, an empty configured token authorizes nobody: this route is +// registered in every deployment, so failing open would publish an +// unauthenticated mux to any caller that can reach the port. +func authorizedPeer(r *http.Request, expected string) bool { + if expected == "" { + return false + } + const prefix = "Bearer " + header := r.Header.Get("Authorization") + if len(header) < len(prefix) || header[:len(prefix)] != prefix { + return false + } + return subtle.ConstantTimeCompare([]byte(header[len(prefix):]), []byte(expected)) == 1 +} diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go new file mode 100644 index 000000000000..b4aa98356c13 --- /dev/null +++ b/core/http/endpoints/cluster/peer_test.go @@ -0,0 +1,117 @@ +package cluster_test + +import ( + "net/http" + "net/http/httptest" + "strings" + + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Peer link handler", func() { + var ( + srv *httptest.Server + sessions chan *yamux.Session + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + e := echo.New() + clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + sessions <- s + }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + wsURL := func(s *httptest.Server) string { + return "ws" + strings.TrimPrefix(s.URL, "http") + "/api/cluster/peer?id=peer-1" + } + + It("rejects a connection with no token", func() { + _, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), nil) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("rejects a connection with the wrong token", func() { + h := http.Header{} + h.Set("Authorization", "Bearer wrong") + _, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("accepts an authenticated peer and yields a usable yamux session", func() { + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + var serverSess *yamux.Session + Eventually(sessions, "5s").Should(Receive(&serverSess)) + Expect(serverSess).ToNot(BeNil()) + + // The client wraps its side as a yamux CLIENT and opens a stream; the + // server must accept it. This proves the WebSocket was adapted into a + // stream-oriented conn correctly, which is the part most likely to be + // subtly wrong. + clientSess, err := yamux.Client(clusterep.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = clientSess.Close() }) + + go func() { + defer GinkgoRecover() + st, e := clientSess.OpenStream(GinkgoT().Context()) + if e == nil { + _, _ = st.Write([]byte("hello")) + } + }() + + accepted := make(chan []byte, 1) + go func() { + defer GinkgoRecover() + st, e := serverSess.AcceptStream() + if e != nil { + return + } + buf := make([]byte, 5) + if _, e := st.Read(buf); e == nil { + accepted <- buf + } + }() + Eventually(accepted, "10s").Should(Receive(Equal([]byte("hello")))) + }) + + It("reports the peer id it was given", func() { + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + ids := make(chan string, 1) + e := echo.New() + clusterep.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + s2 := httptest.NewServer(e) + DeferCleanup(s2.Close) + + conn, _, err := websocket.DefaultDialer.Dial(wsURL(s2), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + Eventually(ids, "5s").Should(Receive(Equal("peer-1"))) + }) +}) + +var _ = Describe("Peer link auth prefix", func() { + It("is covered by the alternative-authentication prefix list", func() { + // /api/cluster/ authenticates with the cluster token, not the global + // session middleware, so it must be listed or every peer dial 401s. + Expect(clusterep.AlternativeAuthPrefix).To(Equal("/api/cluster/")) + }) +}) diff --git a/core/http/endpoints/cluster/wsconn.go b/core/http/endpoints/cluster/wsconn.go new file mode 100644 index 000000000000..1f1f5306ccd2 --- /dev/null +++ b/core/http/endpoints/cluster/wsconn.go @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "fmt" + "io" + "net" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// WebsocketConn adapts a gorilla WebSocket into the net.Conn that a yamux +// session drives. +// +// The two disagree about framing: WebSocket delivers whole messages, yamux +// wants an undelimited byte stream. The adapter therefore keeps the reader of +// the message it is part-way through between calls, so a Read whose buffer is +// smaller than the message hands back a prefix now and the rest next time +// instead of dropping the tail. That case is not hypothetical: yamux reads +// through a 4 KiB bufio.Reader while a single stream write can put a much +// larger data frame on the wire in one Write, so any message above the buffer +// size is read in pieces. +// +// The returned conn is safe for one reader and one writer concurrently, which +// is all yamux uses: its recvLoop reads and its sendLoop writes. It is not a +// general-purpose net.Conn. +func WebsocketConn(ws *websocket.Conn) net.Conn { + return &wsConn{ws: ws} +} + +type wsConn struct { + ws *websocket.Conn + + // readMu guards frame, which carries a partially consumed message across + // Read calls. gorilla allows a single concurrent reader, and this keeps + // the adapter to that contract even if a caller reads from two goroutines. + readMu sync.Mutex + frame io.Reader + + // writeMu keeps to gorilla's one-concurrent-writer contract. + 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) + } + // Binary is the only type this link speaks. Skipping an unexpected + // text message would silently desynchronise the yamux framing, so + // it is reported instead. + if messageType != websocket.BinaryMessage { + return 0, fmt.Errorf("cluster: peer link received websocket message type %d, want binary", messageType) + } + c.frame = r + } + + n, err := c.frame.Read(p) + if err == io.EOF { + // End of one message, not end of the stream: drop the reader so + // the next call pulls the next message. Passing io.EOF up would + // end the yamux session at an arbitrary message boundary. + c.frame = nil + err = nil + } + if n > 0 || err != nil { + return n, err + } + // A zero-length message yields nothing to return, and (0, nil) reads + // look like a stalled stream to some callers, so wait for the next one. + } +} + +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 underlying network connection without negotiating a +// WebSocket close handshake. yamux has already sent its own go-away by this +// point, and a close frame would need the write lock that a blocked sendLoop +// may 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.ws.SetReadDeadline(t); err != nil { + return err + } + return c.ws.SetWriteDeadline(t) +} + +func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) } +func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(t) } + +// translateReadErr maps a peer hanging up cleanly onto io.EOF, which is how a +// yamux session recognises a normal ending. Any other close code, and any +// transport error, is passed through so 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/core/http/endpoints/cluster/wsconn_test.go b/core/http/endpoints/cluster/wsconn_test.go new file mode 100644 index 000000000000..f3ce4ddf0261 --- /dev/null +++ b/core/http/endpoints/cluster/wsconn_test.go @@ -0,0 +1,216 @@ +package cluster_test + +import ( + "bytes" + "crypto/rand" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "time" + + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// wsPair returns the two ends of one live WebSocket connection. +func wsPair() (clientSide, serverSide *websocket.Conn) { + GinkgoHelper() + + upgrader := websocket.Upgrader{} + accepted := make(chan *websocket.Conn, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + accepted <- ws + })) + DeferCleanup(srv.Close) + + c, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = c.Close() }) + + var s *websocket.Conn + Eventually(accepted, "5s").Should(Receive(&s)) + DeferCleanup(func() { _ = s.Close() }) + + return c, s +} + +var _ = Describe("WebsocketConn framing", func() { + // The specs below exist because the brief's end-to-end yamux spec cannot + // catch a lost message tail: yamux reads through a 4 KiB bufio.Reader, so + // every small message arrives whole no matter how the adapter behaves. + // These drive the adapter directly with buffers smaller than the message. + + It("returns the rest of a message on the following Read", func() { + clientWS, serverWS := wsPair() + writer := clusterep.WebsocketConn(clientWS) + reader := clusterep.WebsocketConn(serverWS) + + // A lost tail would otherwise park the reassembly below forever; with a + // deadline it fails as a timeout on the read that has nothing left. + Expect(reader.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + + payload := []byte("0123456789abcdefghijklmnopqrstuvwxyz") + n, err := writer.Write(payload) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(Equal(len(payload))) + + // Deliberately smaller than the message: a naive adapter that starts a + // fresh NextReader on every call drops everything past the first 7 + // bytes, and this reassembly fails. + got := make([]byte, 0, len(payload)) + buf := make([]byte, 7) + for len(got) < len(payload) { + read, err := reader.Read(buf) + Expect(err).ToNot(HaveOccurred()) + Expect(read).To(BeNumerically(">", 0)) + Expect(read).To(BeNumerically("<=", len(buf))) + got = append(got, buf[:read]...) + } + Expect(got).To(Equal(payload)) + }) + + It("streams a message larger than the yamux read buffer without loss or reordering", func() { + clientWS, serverWS := wsPair() + writer := clusterep.WebsocketConn(clientWS) + reader := clusterep.WebsocketConn(serverWS) + + Expect(reader.SetReadDeadline(time.Now().Add(20 * time.Second))).To(Succeed()) + + payload := make([]byte, 256*1024) + _, err := rand.Read(payload) + Expect(err).ToNot(HaveOccurred()) + + go func() { + defer GinkgoRecover() + _, _ = writer.Write(payload) + }() + + // 4096 is the buffer yamux's bufio.Reader actually hands down. + got := make([]byte, len(payload)) + _, err = io.ReadFull(reader, got) + Expect(err).ToNot(HaveOccurred()) + Expect(bytes.Equal(got, payload)).To(BeTrue()) + }) + + It("presents consecutive messages as one continuous byte stream", func() { + clientWS, serverWS := wsPair() + writer := clusterep.WebsocketConn(clientWS) + reader := clusterep.WebsocketConn(serverWS) + + Expect(reader.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + + for _, chunk := range []string{"abc", "", "de", "fghij"} { + _, err := writer.Write([]byte(chunk)) + Expect(err).ToNot(HaveOccurred()) + } + + // A read spanning three messages must be satisfied, and the empty + // message must not surface as a premature (0, nil) or an EOF. + got := make([]byte, 10) + _, err := io.ReadFull(reader, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("abcdefghij")) + }) + + It("reports a clean peer close as io.EOF", func() { + clientWS, serverWS := wsPair() + reader := clusterep.WebsocketConn(serverWS) + + Expect(clientWS.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))).To(Succeed()) + + _, err := reader.Read(make([]byte, 8)) + Expect(err).To(MatchError(io.EOF)) + }) + + It("refuses a text message rather than desynchronising the stream", func() { + clientWS, serverWS := wsPair() + reader := clusterep.WebsocketConn(serverWS) + + Expect(clientWS.WriteMessage(websocket.TextMessage, []byte("not a frame"))).To(Succeed()) + + _, err := reader.Read(make([]byte, 32)) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("want binary")) + }) + + It("satisfies net.Conn, including the deadlines yamux sets on every write", func() { + clientWS, _ := wsPair() + var conn net.Conn = clusterep.WebsocketConn(clientWS) + + Expect(conn.LocalAddr()).ToNot(BeNil()) + Expect(conn.RemoteAddr()).ToNot(BeNil()) + // yamux's sendLoop calls SetWriteDeadline before every flush, so an + // adapter that dropped the call would let a stalled peer block the + // session forever instead of failing it. + Expect(conn.SetWriteDeadline(time.Now().Add(time.Minute))).To(Succeed()) + Expect(conn.SetReadDeadline(time.Now().Add(time.Minute))).To(Succeed()) + Expect(conn.SetDeadline(time.Time{})).To(Succeed()) + }) +}) + +var _ = Describe("Peer link payloads", func() { + It("carries a payload far larger than one yamux frame end to end", func() { + sessions := make(chan *yamux.Session, 1) + e := echo.New() + clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) + srv := httptest.NewServer(e) + DeferCleanup(srv.Close) + + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(srv.URL, "http")+"/api/cluster/peer?id=peer-1", h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + var serverSess *yamux.Session + Eventually(sessions, "5s").Should(Receive(&serverSess)) + + clientSess, err := yamux.Client(clusterep.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = clientSess.Close() }) + + payload := make([]byte, 1<<20) + _, err = rand.Read(payload) + Expect(err).ToNot(HaveOccurred()) + + go func() { + defer GinkgoRecover() + st, e := clientSess.OpenStream(GinkgoT().Context()) + if e != nil { + return + } + defer func() { _ = st.Close() }() + _, _ = io.Copy(st, bytes.NewReader(payload)) + }() + + received := make(chan []byte, 1) + go func() { + defer GinkgoRecover() + st, e := serverSess.AcceptStream() + if e != nil { + return + } + buf := make([]byte, len(payload)) + if _, e := io.ReadFull(st, buf); e == nil { + received <- buf + } + }() + + var got []byte + Eventually(received, "30s").Should(Receive(&got)) + Expect(bytes.Equal(got, payload)).To(BeTrue()) + }) +}) From 5847f6ee8044c70b23201ac5587ed59398659783 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 23:19:25 +0000 Subject: [PATCH 07/42] fix(cluster): stop reading a bare EOF as a clean ending A dead yamux session does not always arrive wrapped. Session.close publishes shutdownErr and closes shutdownCh before it force-closes the streams, so a Write or Close landing in that window gets the raw cause back instead (session.go:507-510, 528-533), and for a peer that vanished the raw cause is a bare io.EOF. The generic io.EOF clause then reported the dead session as a clean completion. Remove the clause. A clean read-side EOF never reached it anyway: io.Copy consumes that and reports nil, and neither *yamux.Stream nor *net.TCPConn takes a WriteTo/ReadFrom path that would hand one back. Every existing spec still passes, the io.EOF entry in the normal-termination table included, which is what showed the branch was dead for legitimate endings and live only for the bug. Add a spec driving a real yamux session end to end. Every mux shape until now was a synthesized error, which is exactly why a race inside the real library stayed invisible. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/splice.go | 31 +++++++--- core/services/cluster/splice_test.go | 88 +++++++++++++++++++++++++++- 2 files changed, 109 insertions(+), 10 deletions(-) diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go index ba5ac28af294..5df97e890da0 100644 --- a/core/services/cluster/splice.go +++ b/core/services/cluster/splice.go @@ -69,15 +69,22 @@ func copyStream(dst io.Writer, src io.Reader) error { } // normalizeStreamErr drops the endings that mean the conversation is over -// rather than broken. io.EOF is the clean end of a stream, net.ErrClosed is -// what a socket reports once it or its peer has been closed, and -// io.ErrClosedPipe is the same condition on an in-memory pipe. +// rather than broken: net.ErrClosed is what a socket reports once it or its +// peer has been closed, and io.ErrClosedPipe is the same condition on an +// in-memory pipe. +// +// io.EOF is deliberately absent. A clean read-side EOF never gets this far, +// because io.Copy consumes it and reports nil, and neither *yamux.Stream nor +// *net.TCPConn takes a WriteTo/ReadFrom path that would hand one back. So a +// bare io.EOF arriving here came from a failing Write or Close, where it means +// the peer is gone, and yamux produces exactly that when a Write races its +// session's shutdown (see isMuxSessionFailure). // // The mux checks run first, and that ordering is load-bearing: a dying yamux -// session hands every live stream its own cause wrapped up (session.go:330), -// and that cause is routinely io.EOF or a closed-socket error, so consulting -// the generic endings first would report a peer that vanished mid-request as a -// clean completion. +// session usually hands every live stream its own cause wrapped up +// (session.go:330), and that cause is routinely a closed-socket error, so +// consulting the generic endings first would report a peer that vanished +// mid-request as a clean completion. func normalizeStreamErr(err error) error { if err == nil { return nil @@ -88,8 +95,7 @@ func normalizeStreamErr(err error) error { if isMuxStreamTeardown(err) { return nil } - if errors.Is(err, io.EOF) || - errors.Is(err, net.ErrClosed) || + if errors.Is(err, net.ErrClosed) || errors.Is(err, io.ErrClosedPipe) { return nil } @@ -122,6 +128,13 @@ func isMuxSessionFailure(err error) bool { // the cause, so the bare sentinel means this stream was reset and a // wrapped one means the session died under it. Identity is what separates // them; errors.Is cannot. + // + // Wrapped is not the only way a dead session shows up, though. close() + // publishes shutdownErr and closes shutdownCh before it force-closes the + // streams, so a Write or Close landing in that window gets the raw cause + // back instead (session.go:507-510, 528-533). That form is unrecognisable + // as yamux at all, which is why normalizeStreamErr no longer forgives a + // bare io.EOF: for a peer that vanished, the raw cause is precisely io.EOF. return errors.Is(err, yamux.ErrStreamReset) && err != yamux.ErrStreamReset } diff --git a/core/services/cluster/splice_test.go b/core/services/cluster/splice_test.go index f18618b6b2cd..7bbe3246de05 100644 --- a/core/services/cluster/splice_test.go +++ b/core/services/cluster/splice_test.go @@ -1,6 +1,7 @@ package cluster_test import ( + "context" "errors" "fmt" "io" @@ -206,6 +207,23 @@ var _ = Describe("Splice", func() { Entry("an internal-error go-away", &yamux.GoAwayError{Remote: true, ErrorCode: 2}), ) + // A bare io.EOF can only reach Splice from a failing Write. io.Copy + // never surfaces a clean read-side EOF, and yamux hands out the raw + // cause rather than the wrapped one when a Write or Close races + // Session.close's shutdown window (session.go:507-510), so for a + // vanished peer this IS the dead session, arriving unwrapped. + It("reports a write that fails with a bare EOF", func() { + sink := &scriptedStream{writeErr: io.EOF} + source := &scriptedStream{feeds: true} + + done := make(chan error, 1) + go func() { done <- cluster.Splice(sink, source) }() + + var err error + Eventually(done, "5s").Should(Receive(&err)) + Expect(err).To(MatchError(io.EOF)) + }) + // The distinction the classifier turns on, in one spec: yamux uses the // same sentinel for "this stream was reset", which Splice provokes // itself and must stay quiet about, and as the head of the wrapped @@ -254,6 +272,57 @@ var _ = Describe("Splice", func() { Eventually(done, "5s").Should(Receive(BeNil())) }) + // Everything above feeds Splice a synthesized error. This one drives a + // real yamux session, because the shapes a live library produces are + // not always the ones its source suggests: the bug this spec was added + // alongside was a race inside Session.close that no synthesized error + // could show. It asserts only that a dead session is reported, not how + // it is spelled, since which of the two forms arrives is a race. + It("reports a real yamux session dying under a live stream", func() { + clientConn, serverConn := net.Pipe() + client, err := yamux.Client(clientConn, nil, nil) + Expect(err).ToNot(HaveOccurred()) + server, err := yamux.Server(serverConn, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { + _ = client.Close() + _ = server.Close() + }) + + accepted := make(chan *yamux.Stream, 1) + go func() { + defer GinkgoRecover() + far, err := server.AcceptStream() + if err != nil { + close(accepted) + return + } + accepted <- far + }() + + stream, err := client.OpenStream(context.Background()) + Expect(err).ToNot(HaveOccurred()) + // Push a byte so the stream is established on both sides before + // the session is killed. + _, err = stream.Write([]byte("x")) + Expect(err).ToNot(HaveOccurred()) + var far *yamux.Stream + Eventually(accepted, "10s").Should(Receive(&far)) + _, err = far.Read(make([]byte, 1)) + Expect(err).ToNot(HaveOccurred()) + + done := make(chan error, 1) + go func() { done <- cluster.Splice(stream, &scriptedStream{}) }() + + // The peer's process disappears: the connection carrying the + // session goes away, which kills every stream riding on it. + Expect(serverConn.Close()).To(Succeed()) + + var spliceErr error + Eventually(done, "10s").Should(Receive(&spliceErr)) + Expect(spliceErr).To(HaveOccurred()) + }) + // The anti-leak guarantee, which the pipe specs cannot see because // their parked copy is released too quickly to catch Splice in the // act. Waking a copy is asynchronous on a real stream (yamux's Close @@ -281,7 +350,11 @@ var _ = Describe("Splice", func() { // parks in Read until Close, standing in for an idle half of a live stream. type scriptedStream struct { readErr error + writeErr error closeErr error + // feeds makes Read produce bytes instead of parking, so a spec can keep a + // direction copying until its destination fails. + feeds bool // holdReadPastClose keeps a parked Read blocked until release is called, // standing in for the gap between a Close waking a reader and that reader // running. Without it, Close releases the Read as a real stream does. @@ -307,11 +380,24 @@ func (s *scriptedStream) Read(p []byte) (int, error) { if s.readErr != nil { return 0, s.readErr } + if s.feeds { + select { + case <-s.gate(): + return 0, io.EOF + default: + return len(p), nil + } + } <-s.gate() return 0, io.EOF } -func (s *scriptedStream) Write(p []byte) (int, error) { return len(p), nil } +func (s *scriptedStream) Write(p []byte) (int, error) { + if s.writeErr != nil { + return 0, s.writeErr + } + return len(p), nil +} func (s *scriptedStream) Close() error { s.closeN.Add(1) From 78958beef1ee30088f3fef941d6e05f9d709c064 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 23:44:11 +0000 Subject: [PATCH 08/42] fix(cluster): own the peer auth prefix in auth, and pin what the specs claimed The prefix constant moves to core/http/auth beside the check that uses it, and the endpoints package derives its route from there. Seven sibling endpoint packages already import auth, so the previous direction would have deadlocked the build as soon as this one registered in RouteFeatureRegistry, and it was dragging echo, gorilla/websocket and yamux into unrelated service packages. Four properties were argued in comments and held by nothing. Flipping the empty-token check to fail open, making SetWriteDeadline a no-op, returning a zero-length read for a zero-length message, and dropping the recover around the callback all left the suite green. Each now fails a spec that asserts the behaviour rather than the setter's return value. SetWriteDeadline takes the write mutex because gorilla keeps that deadline in a plain struct field applied at the next flush; SetReadDeadline must not take the read mutex, since it goes straight to the net.Conn and would otherwise block behind the read it exists to unblock. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/http/auth/public_routes.go | 16 ++-- core/http/endpoints/cluster/peer.go | 28 +++++-- core/http/endpoints/cluster/peer_test.go | 76 ++++++++++++++++++ core/http/endpoints/cluster/wsconn.go | 26 ++++-- core/http/endpoints/cluster/wsconn_test.go | 92 +++++++++++++++++++--- 5 files changed, 211 insertions(+), 27 deletions(-) diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index e1db88a02f60..04d90d07d56c 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -5,8 +5,6 @@ package auth import ( "net/http" "strings" - - clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" ) type publicRouteRule struct { @@ -76,13 +74,17 @@ func isPublicRoute(method, path string) bool { return false } +// ClusterPathPrefix is the replica-to-replica namespace. Its handlers check the +// cluster token in the Authorization header themselves, so the check below lets +// them through the global session middleware. The cluster endpoints build their +// route paths from this same constant, which is why it lives here beside the +// check rather than beside the handlers: the exemption and the route it exempts +// cannot then be changed independently. +const ClusterPathPrefix = "/api/cluster/" + // usesAlternativeAuthentication identifies requests whose credentials are // validated by route-group middleware instead of the global auth middleware. func usesAlternativeAuthentication(path string) bool { - // The peer link carries the cluster token in an Authorization header that - // no browser session ever sets, and its handler checks that token itself. - // The prefix comes from the endpoints package so the route and the - // exemption cannot drift apart. return strings.HasPrefix(path, "/api/node/") || - strings.HasPrefix(path, clusterep.AlternativeAuthPrefix) + strings.HasPrefix(path, ClusterPathPrefix) } diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go index 4b6ef542caf4..5ff916f6c93e 100644 --- a/core/http/endpoints/cluster/peer.go +++ b/core/http/endpoints/cluster/peer.go @@ -9,18 +9,25 @@ package cluster import ( "crypto/subtle" "net/http" + "strings" "github.com/gorilla/websocket" "github.com/labstack/echo/v4" "github.com/libp2p/go-yamux/v5" + "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/xlog" ) // AlternativeAuthPrefix is the path prefix whose credentials are checked by -// this package rather than by the global session middleware. The auth layer -// consults this same constant, so the two cannot drift apart and leave every -// peer dial answering 401. -const AlternativeAuthPrefix = "/api/cluster/" +// this package rather than by the global session middleware. +// +// It is the auth package's own constant rather than a second copy: auth decides +// which paths bypass the session middleware, so owning the prefix there and +// deriving the route from it here means a change to either one moves both. The +// dependency runs endpoints -> auth, the direction the rest of core/http flows; +// pointing it the other way would deadlock the build as soon as this package +// needs anything from auth, which registering in RouteFeatureRegistry will. +const AlternativeAuthPrefix = auth.ClusterPathPrefix // PeerPath is the route a peer replica dials. const PeerPath = AlternativeAuthPrefix + "peer" @@ -80,6 +87,16 @@ func PeerHandler(token string, onSession func(peerID string, sess *yamux.Session } xlog.Debug("cluster peer link established", "peer", peerID, "remote", ws.RemoteAddr().String()) + // net/http recovers a panic from this goroutine but does not close a + // hijacked connection afterwards, so a panicking callback would leave + // the peer holding a link nobody accepts streams on: its opens would + // fill the 256-deep backlog and then hang without an error. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() onSession(peerID, sess) return nil } @@ -95,9 +112,10 @@ func authorizedPeer(r *http.Request, expected string) bool { if expected == "" { return false } + // RFC 7235 makes the scheme case-insensitive; the token after it is not. const prefix = "Bearer " header := r.Header.Get("Authorization") - if len(header) < len(prefix) || header[:len(prefix)] != prefix { + if len(header) < len(prefix) || !strings.EqualFold(header[:len(prefix)], prefix) { return false } return subtle.ConstantTimeCompare([]byte(header[len(prefix):]), []byte(expected)) == 1 diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go index b4aa98356c13..a9f027ff0c56 100644 --- a/core/http/endpoints/cluster/peer_test.go +++ b/core/http/endpoints/cluster/peer_test.go @@ -106,6 +106,82 @@ var _ = Describe("Peer link handler", func() { Eventually(ids, "5s").Should(Receive(Equal("peer-1"))) }) + It("rejects every dial when no cluster token is configured", func() { + // The route is registered in every deployment, so an empty configured + // token must authorize nobody. Failing open the way the worker + // file-transfer server's checkBearerToken does would publish an + // unauthenticated yamux multiplexer to anyone who can reach the port. + e := echo.New() + accepted := make(chan *yamux.Session, 1) + clusterep.RegisterClusterRoutes(e, "", func(_ string, sess *yamux.Session) { accepted <- sess }) + s2 := httptest.NewServer(e) + DeferCleanup(s2.Close) + + for _, header := range []http.Header{nil, {"Authorization": []string{"Bearer "}}, {"Authorization": []string{"Bearer anything"}}} { + _, resp, err := websocket.DefaultDialer.Dial(wsURL(s2), header) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + } + Expect(accepted).ToNot(Receive()) + }) + + It("accepts the bearer scheme in any case", func() { + // RFC 7235 makes the scheme case-insensitive. The token after it is not. + h := http.Header{} + h.Set("Authorization", "bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Eventually(sessions, "5s").Should(Receive()) + }) + + It("closes the session when the callback panics", func() { + // net/http recovers the panic but leaves the hijacked socket open, so + // without the handler's own recover the peer would keep a link nobody + // ever accepts streams on. + e := echo.New() + clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, _ *yamux.Session) { + panic("callback exploded") + }) + s2 := httptest.NewServer(e) + DeferCleanup(func() { + // A hijacked connection the handler never closed would park + // httptest's Close forever, turning the assertion below into a + // suite hang. Forcing the conns shut keeps the failure legible. + s2.CloseClientConnections() + s2.Close() + }) + + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsURL(s2), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + clientSess, err := yamux.Client(clusterep.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = clientSess.Close() }) + + // Asserting on OpenStream would hang rather than fail: yamux only + // acknowledges a stream once the peer accepts it, and the leak this + // pins is precisely that nobody ever will. The session's own liveness + // is the observable that answers in both directions. + Eventually(clientSess.IsClosed, "10s").Should(BeTrue()) + }) + + It("rejects an authenticated dial that names no peer", func() { + // The session is keyed by peer id, so a nameless link could never be + // looked up again; refusing it is cheaper than leaking it. + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + _, resp, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(srv.URL, "http")+"/api/cluster/peer", h) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + Expect(sessions).ToNot(Receive()) + }) }) var _ = Describe("Peer link auth prefix", func() { diff --git a/core/http/endpoints/cluster/wsconn.go b/core/http/endpoints/cluster/wsconn.go index 1f1f5306ccd2..ca932d741cc8 100644 --- a/core/http/endpoints/cluster/wsconn.go +++ b/core/http/endpoints/cluster/wsconn.go @@ -24,8 +24,9 @@ import ( // larger data frame on the wire in one Write, so any message above the buffer // size is read in pieces. // -// The returned conn is safe for one reader and one writer concurrently, which -// is all yamux uses: its recvLoop reads and its sendLoop writes. It is not a +// The returned conn is safe for one reader and one writer concurrently, plus a +// third goroutine setting deadlines, which is what the relay needs: yamux's +// sendLoop writes while a supervisor arms an idle deadline. It is not a // general-purpose net.Conn. func WebsocketConn(ws *websocket.Conn) net.Conn { return &wsConn{ws: ws} @@ -105,14 +106,27 @@ 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.ws.SetReadDeadline(t); err != nil { + if err := c.SetReadDeadline(t); err != nil { return err } - return c.ws.SetWriteDeadline(t) + return c.SetWriteDeadline(t) } -func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) } -func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(t) } +// SetReadDeadline needs no lock, and must not take readMu: gorilla passes the +// read deadline straight to the underlying net.Conn, whose deadline setters are +// safe to call from another goroutine, and taking readMu would block behind the +// parked Read this call exists to unblock. +func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) } + +// SetWriteDeadline takes writeMu because gorilla stores the write deadline in a +// plain struct field (conn.go:787) and applies it when it next flushes, so +// setting it while a write is in flight is a data race, not merely a late bound. +func (c *wsConn) SetWriteDeadline(t time.Time) error { + c.writeMu.Lock() + defer c.writeMu.Unlock() + + return c.ws.SetWriteDeadline(t) +} // translateReadErr maps a peer hanging up cleanly onto io.EOF, which is how a // yamux session recognises a normal ending. Any other close code, and any diff --git a/core/http/endpoints/cluster/wsconn_test.go b/core/http/endpoints/cluster/wsconn_test.go index f3ce4ddf0261..5ef3fcf3b1e2 100644 --- a/core/http/endpoints/cluster/wsconn_test.go +++ b/core/http/endpoints/cluster/wsconn_test.go @@ -7,6 +7,7 @@ import ( "net" "net/http" "net/http/httptest" + "os" "strings" "time" @@ -115,14 +116,34 @@ var _ = Describe("WebsocketConn framing", func() { Expect(err).ToNot(HaveOccurred()) } - // A read spanning three messages must be satisfied, and the empty - // message must not surface as a premature (0, nil) or an EOF. + // A read spanning several messages must be satisfied: a message + // boundary is not the end of the stream. got := make([]byte, 10) _, err := io.ReadFull(reader, got) Expect(err).ToNot(HaveOccurred()) Expect(string(got)).To(Equal("abcdefghij")) }) + It("never hands back a zero-length read for a zero-length message", func() { + clientWS, serverWS := wsPair() + writer := clusterep.WebsocketConn(clientWS) + reader := clusterep.WebsocketConn(serverWS) + + Expect(reader.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + + // An empty message carries nothing to return. Handing back (0, nil) + // would be legal for io.Reader but reads as a stalled stream to callers + // that loop on n, so the adapter waits for the next message instead. + _, err := writer.Write(nil) + Expect(err).ToNot(HaveOccurred()) + _, err = writer.Write([]byte("xy")) + Expect(err).ToNot(HaveOccurred()) + + n, err := reader.Read(make([]byte, 8)) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(Equal(2)) + }) + It("reports a clean peer close as io.EOF", func() { clientWS, serverWS := wsPair() reader := clusterep.WebsocketConn(serverWS) @@ -145,18 +166,71 @@ var _ = Describe("WebsocketConn framing", func() { Expect(err.Error()).To(ContainSubstring("want binary")) }) - It("satisfies net.Conn, including the deadlines yamux sets on every write", func() { + It("satisfies net.Conn", func() { clientWS, _ := wsPair() var conn net.Conn = clusterep.WebsocketConn(clientWS) Expect(conn.LocalAddr()).ToNot(BeNil()) Expect(conn.RemoteAddr()).ToNot(BeNil()) - // yamux's sendLoop calls SetWriteDeadline before every flush, so an - // adapter that dropped the call would let a stalled peer block the - // session forever instead of failing it. - Expect(conn.SetWriteDeadline(time.Now().Add(time.Minute))).To(Succeed()) - Expect(conn.SetReadDeadline(time.Now().Add(time.Minute))).To(Succeed()) - Expect(conn.SetDeadline(time.Time{})).To(Succeed()) + }) + + It("enforces a write deadline, which yamux arms before every flush", func() { + clientWS, _ := wsPair() + conn := clusterep.WebsocketConn(clientWS) + + // Asserting that the setter returns nil would prove nothing: gorilla + // only records the deadline and applies it at the next flush. The write + // below is what shows the deadline reached the socket, and an adapter + // that swallowed the call would let a stalled peer block yamux's send + // loop forever instead of failing it. + Expect(conn.SetWriteDeadline(time.Now().Add(-time.Second))).To(Succeed()) + _, err := conn.Write([]byte("x")) + Expect(err).To(HaveOccurred()) + Expect(os.IsTimeout(err)).To(BeTrue(), "want a timeout, got %v", err) + }) + + It("enforces a read deadline, which is how a parked reader is unblocked", func() { + clientWS, _ := wsPair() + conn := clusterep.WebsocketConn(clientWS) + + Expect(conn.SetReadDeadline(time.Now().Add(-time.Second))).To(Succeed()) + _, err := conn.Read(make([]byte, 8)) + Expect(err).To(HaveOccurred()) + Expect(os.IsTimeout(err)).To(BeTrue(), "want a timeout, got %v", err) + }) + + It("arms both directions from SetDeadline", func() { + clientWS, _ := wsPair() + conn := clusterep.WebsocketConn(clientWS) + + Expect(conn.SetDeadline(time.Now().Add(-time.Second))).To(Succeed()) + + _, err := conn.Read(make([]byte, 8)) + Expect(os.IsTimeout(err)).To(BeTrue(), "read: want a timeout, got %v", err) + _, err = conn.Write([]byte("x")) + Expect(os.IsTimeout(err)).To(BeTrue(), "write: want a timeout, got %v", err) + }) + + It("lets a deadline be armed while another goroutine writes", func() { + // Task 5's relay arms an idle deadline from a supervisor goroutine while + // yamux's send loop writes. gorilla keeps the write deadline in a plain + // struct field, so this is a data race unless the adapter serialises it; + // the spec is here to be run under -race, where it would report one. + clientWS, _ := wsPair() + conn := clusterep.WebsocketConn(clientWS) + + done := make(chan struct{}) + go func() { + defer GinkgoRecover() + defer close(done) + for i := 0; i < 200; i++ { + _, _ = conn.Write([]byte("ping")) + } + }() + for i := 0; i < 200; i++ { + _ = conn.SetWriteDeadline(time.Now().Add(time.Minute)) + } + Eventually(done, "20s").Should(BeClosed()) }) }) From e7aac9b52beca28109acf1b02d8a128766428925 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 00:09:15 +0000 Subject: [PATCH 09/42] feat(cluster): dial and pool yamux links to peer replicas Distinguishes a peer missing from the registry from a peer that will not answer: the second must never be readable as node absence, or a network hiccup between replicas evicts healthy workers. The distinction is a property of the error type rather than of the call sites. The unreachable error formats its cause into its message and keeps it out of its unwrap chain, so an ErrInstanceNotFound picked up on the dial path cannot reach a caller's absence check. One yamux session is cached per peer and re-dialled when OpenStream on it fails, which is how both a dead transport and a graceful remote go-away arrive. A reset of one stream never reaches the pool, so an abandoned request cannot cost every other worker its link. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/peerlink.go | 302 ++++++++++++++++++ .../cluster/peerlink_internal_test.go | 41 +++ core/services/cluster/peerlink_test.go | 221 +++++++++++++ 3 files changed, 564 insertions(+) create mode 100644 core/services/cluster/peerlink.go create mode 100644 core/services/cluster/peerlink_internal_test.go create mode 100644 core/services/cluster/peerlink_test.go diff --git a/core/services/cluster/peerlink.go b/core/services/cluster/peerlink.go new file mode 100644 index 000000000000..d4a21b4c59b7 --- /dev/null +++ b/core/services/cluster/peerlink.go @@ -0,0 +1,302 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "net/url" + "sync" + "time" + + "github.com/gorilla/websocket" + "github.com/libp2p/go-yamux/v5" + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/xlog" +) + +// ErrPeerUnreachable reports that a peer this deployment knows about could not +// be reached: the dial failed, the peer refused the credentials, or its +// multiplexer would not carry a stream. +// +// It is deliberately NOT a form of ErrInstanceNotFound, and the two must stay +// unmixable. A caller that sees absence is entitled to conclude a node is gone +// and reclaim what it was running; a caller that sees unreachability may only +// retry. Collapsing the two means a network hiccup between two healthy +// replicas evicts healthy workers. See unreachableError for how that is +// enforced rather than merely documented. +var ErrPeerUnreachable = errors.New("cluster: peer unreachable") + +// unreachableError reports a peer that could not be reached, keeping the +// underlying cause in its message and out of its unwrap chain. +// +// Withholding the cause from errors.Is is the point. The dial path resolves a +// peer's address through the registry, so ErrInstanceNotFound is a cause this +// error can genuinely be built over: a row deleted between two attempts, for +// one. If the cause were unwrapped, that failure would satisfy both sentinels +// at once and every caller's absence check would fire on a transport problem. +// The guarantee therefore belongs to the type: no call site can leak absence +// through it, because there is no path by which absence gets out. +type unreachableError struct { + peerID string + cause error +} + +func (e *unreachableError) Error() string { + return fmt.Sprintf("cluster: peer %q unreachable: %v", e.peerID, e.cause) +} + +// Unwrap reports only ErrPeerUnreachable. The cause reaches a human through +// Error() and reaches no error-matching caller at all. +func (e *unreachableError) Unwrap() error { return ErrPeerUnreachable } + +func unreachablePeer(peerID string, cause error) error { + return &unreachableError{peerID: peerID, cause: cause} +} + +// ErrPoolClosed reports an Open on a pool that has been shut down. It is a +// third condition on purpose: the pool being closed is a fact about this +// process and says nothing about whether the peer exists or answers. +var ErrPoolClosed = errors.New("cluster: peer pool is closed") + +const ( + // peerLinkHandshakeTimeout bounds the WebSocket upgrade. It also bounds + // how long Close can wait behind an in-flight dial, since a dial holds the + // per-peer lock Close needs to reach the cached session. + peerLinkHandshakeTimeout = 10 * time.Second + + // peerLinkInitialWindow is the per-stream receive window every stream on a + // peer link starts at, raised from yamux's 256 KiB default. + // + // yamux already bounds head-of-line blocking with MaxMessageSize (64 KiB + // by default), so one stream cannot monopolise the connection whatever the + // window is. What the small default costs is the ramp: a stream carrying a + // multi-megabyte gRPC message spends its first megabytes window-parked, + // paying a round trip per doubling (stream.go:229) before it reaches full + // rate. On a link that is also carrying token streams for other workers, + // that ramp is pure added latency on the bulk transfer for no benefit. + peerLinkInitialWindow = 4 * 1024 * 1024 + + // peerLinkMaxWindow is the ceiling the auto-tuner may grow a stream to, + // raised from yamux's 16 MiB default to cover the bandwidth-delay product + // of a fast cross-zone link (roughly 31 MiB at 10 Gbps and 25 ms). + // + // The window is a cap on data received but not yet read, so the worst case + // a peer can make this replica buffer is MaxIncomingStreams times this, + // which is why MaxIncomingStreams is left at yamux's default rather than + // raised alongside it. + peerLinkMaxWindow = 32 * 1024 * 1024 +) + +// peerLinkConfig returns the yamux configuration for a replica-to-replica link. +func peerLinkConfig() *yamux.Config { + cfg := yamux.DefaultConfig() + cfg.InitialStreamWindowSize = peerLinkInitialWindow + cfg.MaxStreamWindowSize = peerLinkMaxWindow + return cfg +} + +// PeerPool dials peer replicas and keeps one multiplexed session per peer. +// +// A peer link carries traffic for every worker that peer owns, so it is pooled +// rather than dialled per request: a dial per relayed request would add a +// WebSocket handshake to every inference. +// +// The pool needs no knowledge of yamux error shapes to keep its cache honest. +// The two conditions worth reacting to arrive as OpenStream failures and are +// handled by the same retry: a peer that shut down gracefully hands its +// session ErrRemoteGoAway and closes it, and a session whose transport died +// hands out its shutdown error. Conditions scoped to a single stream, such as +// a peer resetting one request, never reach the pool at all, which is right: +// dropping the session over one reset request would tear down every other +// worker's traffic on that link. +type PeerPool struct { + selfID string + token string + reg *Registry + + dialer *websocket.Dialer + + mu sync.Mutex + links map[string]*peerLink + closed bool +} + +// peerLink is the cached session for one peer, plus the lock that serialises +// dialling it. The lock is per-peer so a slow or hanging dial to one peer does +// not hold up opens to any other. +type peerLink struct { + mu sync.Mutex + sess *yamux.Session +} + +// NewPeerPool returns a pool that dials peers as selfID, authenticating with +// the deployment's cluster token. +func NewPeerPool(selfID, token string, reg *Registry) *PeerPool { + return &PeerPool{ + selfID: selfID, + token: token, + reg: reg, + dialer: &websocket.Dialer{ + HandshakeTimeout: peerLinkHandshakeTimeout, + // No Proxy: a peer link is replica-to-replica inside one + // deployment, and honouring HTTP_PROXY would route it through + // whatever egress proxy the environment happens to name. + }, + links: map[string]*peerLink{}, + } +} + +// Open returns a stream to peerID, dialling and caching the session on first +// use. +// +// The errors are three distinct conditions and callers act differently on +// them: ErrInstanceNotFound means the peer is not part of this deployment, +// ErrPeerUnreachable means it is but will not answer, and ErrPoolClosed means +// this process is shutting down. Only the first is node 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 context expired must not cost every other worker + // its link: the session is fine, this request is not. + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + // A session that died between calls is the common case, not an + // exception, so this is a debug line and not a warning. + xlog.Debug("cluster peer link session unusable, re-dialling", "peer", peerID, "error", err) + _ = l.sess.Close() + l.sess = nil + } + + sess, err := p.dial(ctx, peerID) + if err != nil { + return nil, err + } + + st, err := sess.OpenStream(ctx) + if err != nil { + // The peer answered and completed a handshake but will not carry a + // stream, which is a transport condition and never absence. + _ = sess.Close() + return nil, unreachablePeer(peerID, err) + } + + l.sess = sess + return st, nil +} + +// link returns the per-peer entry, creating it on first use. +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 peer's advertised address and brings up one yamux client +// session over an authenticated WebSocket. +// +// A registry miss is returned unchanged so ErrInstanceNotFound reaches the +// caller; everything after it is wrapped as unreachable. +func (p *PeerPool) dial(ctx context.Context, peerID string) (*yamux.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: replica-to-replica TLS is not part of this phase, and the + // link is authenticated by the cluster token rather than by transport. + Scheme: "ws", + Host: inst.AdvertisedAddr, + Path: clusterep.PeerPath, + RawQuery: url.Values{"id": []string{p.selfID}}.Encode(), + } + header := http.Header{} + header.Set("Authorization", "Bearer "+p.token) + + ws, resp, err := p.dialer.DialContext(ctx, endpoint.String(), header) + if resp != nil && resp.Body != nil { + // gorilla hands back the failed handshake's response so a caller can + // read the status; nothing here needs the body, but it has to be + // drained or the connection is not returned to the transport. + _ = resp.Body.Close() + } + if err != nil { + return nil, unreachablePeer(peerID, err) + } + + // Client side of the mux: the dialling replica owns the odd stream IDs, + // matching the yamux.Server the peer handler puts on its end. + sess, err := yamux.Client(clusterep.WebsocketConn(ws), peerLinkConfig(), nil) + 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("cluster peer link dialled", "peer", peerID, "addr", inst.AdvertisedAddr) + return sess, nil +} + +// Close closes every cached session. It is safe to call twice, and an Open +// after it reports ErrPoolClosed rather than anything a caller could read as +// node absence. +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 is closed under its own peer lock rather than under p.mu, + // so closing the pool cannot deadlock against an Open that is mid-dial and + // about to take p.mu to re-check p.closed. + 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/cluster/peerlink_internal_test.go b/core/services/cluster/peerlink_internal_test.go new file mode 100644 index 000000000000..e61a4a4bd2de --- /dev/null +++ b/core/services/cluster/peerlink_internal_test.go @@ -0,0 +1,41 @@ +package cluster + +// These specs are in-package because the property they pin is a property of +// the error TYPE, not of any call site. Asserting it only from outside would +// re-check the paths peerlink_test.go already drives, which leaves the type +// free to start leaking its cause the moment a new call site is added. +// +// The other direction of the rule (a node absent from the registry is not +// merely unreachable) is driven end to end by peerlink_test.go through the +// real Registry, so it is not restated here. + +import ( + "errors" + "fmt" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Peer unreachability is not node absence", func() { + It("stays a transport error even when its cause is an absence error", func() { + // The dial path resolves the peer's address through the registry, so + // an ErrInstanceNotFound is genuinely reachable as a dial cause (a row + // deleted between the lookup and a retry, say). If the type let that + // through, a peer that merely would not answer would read as an absent + // node, and a replica acting on absence evicts healthy workers. + err := unreachablePeer("peer-1", fmt.Errorf("resolving: %w", ErrInstanceNotFound)) + + Expect(errors.Is(err, ErrPeerUnreachable)).To(BeTrue()) + Expect(errors.Is(err, ErrInstanceNotFound)).To(BeFalse(), + "the unreachable error must not unwrap to its cause, or absence leaks through it") + }) + + It("keeps the cause legible in its message", func() { + // Withholding the cause from errors.Is must not withhold it from a + // human reading a log line. + err := unreachablePeer("peer-1", errors.New("connection refused")) + Expect(err.Error()).To(ContainSubstring("peer-1")) + Expect(err.Error()).To(ContainSubstring("connection refused")) + }) +}) diff --git a/core/services/cluster/peerlink_test.go b/core/services/cluster/peerlink_test.go new file mode 100644 index 000000000000..4b145194af06 --- /dev/null +++ b/core/services/cluster/peerlink_test.go @@ -0,0 +1,221 @@ +package cluster_test + +import ( + "context" + "io" + "net/http/httptest" + "strings" + + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("Peer pool", func() { + var ( + db *gorm.DB + reg *cluster.Registry + pool *cluster.PeerPool + srv *httptest.Server + accepted chan *yamux.Session + ctx context.Context + ) + + // startPeer stands up a real peer server and registers it under peerID. + startPeer := func(peerID string) *httptest.Server { + e := echo.New() + clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + accepted <- s + }) + ts := httptest.NewServer(e) + addr := strings.TrimPrefix(ts.URL, "http://") + Expect(reg.Register(ctx, peerID, addr, "test")).To(Succeed()) + return ts + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) + reg = cluster.NewRegistry(db) + accepted = make(chan *yamux.Session, 4) + pool = cluster.NewPeerPool("self", "peer-token", reg) + DeferCleanup(pool.Close) + srv = startPeer("peer-1") + DeferCleanup(srv.Close) + }) + + It("opens a working stream to a live peer", func() { + st, err := pool.Open(ctx, "peer-1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + + var serverSess *yamux.Session + Eventually(accepted, "10s").Should(Receive(&serverSess)) + + go func() { + defer GinkgoRecover() + _, _ = st.Write([]byte("ping")) + }() + + got := make(chan []byte, 1) + go func() { + defer GinkgoRecover() + in, e := serverSess.AcceptStream() + if e != nil { + return + } + buf := make([]byte, 4) + if _, e := io.ReadFull(in, buf); e == nil { + got <- buf + } + }() + Eventually(got, "10s").Should(Receive(Equal([]byte("ping")))) + }) + + It("identifies itself to the peer by its own instance id", func() { + // The peer records which replica is on the far end of the link, so a + // pool that sent the peer's id (or nothing) would leave every inbound + // link anonymous and indistinguishable from every other. + ids := make(chan string, 1) + e := echo.New() + clusterep.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + ts := httptest.NewServer(e) + DeferCleanup(ts.Close) + Expect(reg.Register(ctx, "peer-named", strings.TrimPrefix(ts.URL, "http://"), "test")).To(Succeed()) + + st, err := pool.Open(ctx, "peer-named") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + Eventually(ids, "10s").Should(Receive(Equal("self"))) + }) + + It("reuses one session across opens rather than dialling per stream", func() { + a, err := pool.Open(ctx, "peer-1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = a.Close() }) + Eventually(accepted, "10s").Should(Receive()) + + b, err := pool.Open(ctx, "peer-1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = b.Close() }) + + // A second dial would deliver a second server session. One session + // serving both streams is the property under test: peer links are + // pooled, not per-stream. + Consistently(accepted, "2s", "200ms").ShouldNot(Receive()) + }) + + It("returns ErrPeerUnreachable when the peer is registered but not listening", func() { + dead := startPeer("peer-dead") + dead.Close() + + _, err := pool.Open(ctx, "peer-dead") + Expect(err).To(MatchError(cluster.ErrPeerUnreachable), + "a peer that will not answer must be a transport error, never node absence") + Expect(err).ToNot(MatchError(cluster.ErrInstanceNotFound), + "an unreachable peer must never be readable as an absent node; a replica acting on absence evicts healthy workers") + }) + + It("returns ErrPeerUnreachable when the peer answers but rejects the credentials", func() { + // A token mismatch is a live peer refusing the link, not a missing + // row. Reporting absence here would evict every worker behind a peer + // that was merely rolled out with a stale secret. + e := echo.New() + clusterep.RegisterClusterRoutes(e, "a-different-token", func(_ string, s *yamux.Session) { accepted <- s }) + ts := httptest.NewServer(e) + DeferCleanup(ts.Close) + Expect(reg.Register(ctx, "peer-strict", strings.TrimPrefix(ts.URL, "http://"), "test")).To(Succeed()) + + _, err := pool.Open(ctx, "peer-strict") + Expect(err).To(MatchError(cluster.ErrPeerUnreachable)) + Expect(err).ToNot(MatchError(cluster.ErrInstanceNotFound)) + }) + + It("returns ErrInstanceNotFound when the peer is not in the registry", func() { + _, err := pool.Open(ctx, "never-registered") + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + Expect(err).ToNot(MatchError(cluster.ErrPeerUnreachable), + "a node that was never registered is absent, not merely unreachable") + }) + + It("re-dials after the cached session dies", func() { + first, err := pool.Open(ctx, "peer-1") + Expect(err).ToNot(HaveOccurred()) + Expect(first.Close()).To(Succeed()) + + var serverSess *yamux.Session + Eventually(accepted, "10s").Should(Receive(&serverSess)) + Expect(serverSess.Close()).To(Succeed()) + srv.Close() + + // A replacement peer comes back on a new address under the same id, + // which is what a restarted replica looks like. + replacement := startPeer("peer-1") + DeferCleanup(replacement.Close) + + Eventually(func() error { + st, e := pool.Open(ctx, "peer-1") + if e == nil { + _ = st.Close() + } + return e + }, "15s", "500ms").Should(Succeed()) + + // The replacement's own session proves the pool re-dialled the address + // it re-read from the registry rather than resurrecting the dead one. + Eventually(accepted, "10s").Should(Receive()) + }) + + It("does not drop the pooled session when a single stream is reset by the peer", func() { + // A peer-initiated stream reset is scoped to one request. Dropping the + // session on it would tear down every other worker's traffic on the + // same link, so the pool must keep the session and hand out a fresh + // stream on it. + st, err := pool.Open(ctx, "peer-1") + Expect(err).ToNot(HaveOccurred()) + + var serverSess *yamux.Session + Eventually(accepted, "10s").Should(Receive(&serverSess)) + + go func() { + defer GinkgoRecover() + _, _ = st.Write([]byte("x")) + }() + var inbound *yamux.Stream + Eventually(func() error { + s, e := serverSess.AcceptStream() + inbound = s + return e + }, "10s").Should(Succeed()) + // Reset, not a graceful close: this is the *StreamError{Remote:true} + // the far end sends when it abandons a request. + Expect(inbound.Reset()).To(Succeed()) + Eventually(func() error { + _, e := st.Write([]byte("y")) + return e + }, "10s", "100ms").Should(HaveOccurred()) + + next, err := pool.Open(ctx, "peer-1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = next.Close() }) + Consistently(accepted, "2s", "200ms").ShouldNot(Receive(), + "a reset stream must not cost the whole peer link") + }) + + It("refuses to open after Close and is safe to close twice", func() { + pool.Close() + pool.Close() + + _, err := pool.Open(ctx, "peer-1") + Expect(err).To(HaveOccurred()) + Expect(err).ToNot(MatchError(cluster.ErrInstanceNotFound), + "a locally closed pool says nothing about whether the node exists") + }) +}) From d73730b545bcc222516a19de2afabf58481b48a4 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 00:29:23 +0000 Subject: [PATCH 10/42] refactor(cluster): make the cluster service a leaf and blame the caller's deadline The peer link's WebSocket adapter and route constant lived in core/http/endpoints/cluster, so the dialler in core/services/cluster had to import an HTTP endpoints package to reach them. That pulled echo, core/http/auth and core/config into a package whose doc says it is deliberately free of such dependencies, and it made core/services/nodes reach an endpoints package transitively. It also has no way forward: the worker-connect handler needs the tunnel registry and the node token store, both of which are cycles from there. Move WebsocketConn and PeerPath into core/services/cluster and let the endpoints package import it, which is the direction the rest of core/http flows. The route and the auth exemption still cannot drift apart, now asserted where both are visible rather than by a const reference across the boundary, and the assertion is stronger than the one it replaces: it pins the route under the prefix instead of pinning the prefix's spelling. Also guard the fresh-dial path with ctx.Err(), mirroring the cached path. A caller with a 300ms deadline dialling a live, listening peer was told the peer was unreachable, which would be enough for one impatient client to get a healthy replica routed around once the relay consults these errors. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/cluster/peer.go | 26 +++---- core/http/endpoints/cluster/peer_test.go | 21 ++++-- core/services/cluster/peerlink.go | 56 ++++++++++++-- core/services/cluster/peerlink_test.go | 74 +++++++++++++++++++ .../endpoints => services}/cluster/wsconn.go | 0 .../cluster/wsconn_test.go | 33 +++++---- 6 files changed, 165 insertions(+), 45 deletions(-) rename core/{http/endpoints => services}/cluster/wsconn.go (100%) rename core/{http/endpoints => services}/cluster/wsconn_test.go (91%) diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go index 5ff916f6c93e..85985eee5762 100644 --- a/core/http/endpoints/cluster/peer.go +++ b/core/http/endpoints/cluster/peer.go @@ -14,28 +14,20 @@ import ( "github.com/gorilla/websocket" "github.com/labstack/echo/v4" "github.com/libp2p/go-yamux/v5" - "github.com/mudler/LocalAI/core/http/auth" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/xlog" ) -// AlternativeAuthPrefix is the path prefix whose credentials are checked by -// this package rather than by the global session middleware. -// -// It is the auth package's own constant rather than a second copy: auth decides -// which paths bypass the session middleware, so owning the prefix there and -// deriving the route from it here means a change to either one moves both. The -// dependency runs endpoints -> auth, the direction the rest of core/http flows; -// pointing it the other way would deadlock the build as soon as this package -// needs anything from auth, which registering in RouteFeatureRegistry will. -const AlternativeAuthPrefix = auth.ClusterPathPrefix - -// PeerPath is the route a peer replica dials. -const PeerPath = AlternativeAuthPrefix + "peer" - // RegisterClusterRoutes registers the peer link. onPeer receives every // authenticated session; see PeerHandler for what it is expected to do with it. +// +// The route is core/services/cluster's own constant, so the handler and the +// dialler cannot be registered and dialled at different paths. That the path +// also falls under auth.ClusterPathPrefix, and so bypasses the session +// middleware, is asserted by a spec in this package: it is the only place that +// can see both, since core/services/cluster must not import core/http/auth. func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { - e.GET(PeerPath, PeerHandler(token, onPeer)) + e.GET(clustersvc.PeerPath, PeerHandler(token, onPeer)) } // PeerHandler upgrades an authenticated peer dial to a WebSocket, wraps it as @@ -72,7 +64,7 @@ func PeerHandler(token string, onSession func(peerID string, sess *yamux.Session // Server side of the mux: the dialing peer is the client, so it owns // the odd stream IDs and this side the even ones. - sess, err := yamux.Server(WebsocketConn(ws), nil, nil) + sess, err := yamux.Server(clustersvc.WebsocketConn(ws), nil, nil) if err != nil { xlog.Error("cluster peer link session setup failed", "peer", peerID, "error", err) _ = ws.Close() diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go index a9f027ff0c56..a86ab8add8bd 100644 --- a/core/http/endpoints/cluster/peer_test.go +++ b/core/http/endpoints/cluster/peer_test.go @@ -5,7 +5,9 @@ import ( "net/http/httptest" "strings" + "github.com/mudler/LocalAI/core/http/auth" clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" "github.com/gorilla/websocket" "github.com/labstack/echo/v4" @@ -64,7 +66,7 @@ var _ = Describe("Peer link handler", func() { // server must accept it. This proves the WebSocket was adapted into a // stream-oriented conn correctly, which is the part most likely to be // subtly wrong. - clientSess, err := yamux.Client(clusterep.WebsocketConn(conn), nil, nil) + clientSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = clientSess.Close() }) @@ -159,7 +161,7 @@ var _ = Describe("Peer link handler", func() { Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) - clientSess, err := yamux.Client(clusterep.WebsocketConn(conn), nil, nil) + clientSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = clientSess.Close() }) @@ -185,9 +187,16 @@ var _ = Describe("Peer link handler", func() { }) var _ = Describe("Peer link auth prefix", func() { - It("is covered by the alternative-authentication prefix list", func() { - // /api/cluster/ authenticates with the cluster token, not the global - // session middleware, so it must be listed or every peer dial 401s. - Expect(clusterep.AlternativeAuthPrefix).To(Equal("/api/cluster/")) + It("keeps the peer route inside the alternative-authentication prefix", func() { + // The peer route authenticates with the cluster token, not the global + // session middleware, which only holds while the route sits under the + // prefix auth exempts. Moving either one alone 401s every peer dial. + // + // This lives here because it is the only package that can see both: + // core/services/cluster owns the route and must stay free of any + // core/http dependency, and core/http/auth owns the exemption. + Expect(strings.HasPrefix(clustersvc.PeerPath, auth.ClusterPathPrefix)).To(BeTrue(), + "peer route %q is no longer under the auth-exempt prefix %q", + clustersvc.PeerPath, auth.ClusterPathPrefix) }) }) diff --git a/core/services/cluster/peerlink.go b/core/services/cluster/peerlink.go index d4a21b4c59b7..961c1a3fa3d9 100644 --- a/core/services/cluster/peerlink.go +++ b/core/services/cluster/peerlink.go @@ -12,10 +12,20 @@ import ( "github.com/gorilla/websocket" "github.com/libp2p/go-yamux/v5" - clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" "github.com/mudler/xlog" ) +// PeerPath is the route a replica dials to open a peer link, and the route the +// HTTP layer registers the handler on. It lives here, with the dialler and the +// WebSocket adapter, so that core/services/cluster stays a leaf: the HTTP +// endpoints package imports this one, never the other way round. +// +// The literal is spelled out rather than derived from auth.ClusterPathPrefix +// because importing core/http/auth is exactly the dependency this package must +// not have. The two are kept from drifting apart by a spec in the endpoints +// package, which can see both. +const PeerPath = "/api/cluster/peer" + // ErrPeerUnreachable reports that a peer this deployment knows about could not // be reached: the dial failed, the peer refused the credentials, or its // multiplexer would not carry a stream. @@ -83,9 +93,13 @@ const ( // of a fast cross-zone link (roughly 31 MiB at 10 Gbps and 25 ms). // // The window is a cap on data received but not yet read, so the worst case - // a peer can make this replica buffer is MaxIncomingStreams times this, - // which is why MaxIncomingStreams is left at yamux's default rather than - // raised alongside it. + // a peer can make this replica buffer is MaxIncomingStreams times this. + // At yamux's default MaxIncomingStreams of 1000 that ceiling goes from + // about 15.6 GiB to about 31 GiB per peer session, which is the figure to + // size a replica against; it is why MaxIncomingStreams is left at the + // default rather than raised alongside the window. Both are ceilings on + // unread data and not allocations: yamux grows a stream's receive buffer + // as data arrives. peerLinkMaxWindow = 32 * 1024 * 1024 ) @@ -183,6 +197,17 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { sess, err := p.dial(ctx, peerID) if err != nil { + // Same rule as above, on the path that has no cached session to + // protect: a dial that ran out of the caller's time says nothing about + // the peer, which may be listening and perfectly healthy. Blaming it + // would let one impatient client get a good replica routed around. + // + // This also swallows a genuine ErrInstanceNotFound when the context + // happened to expire at the same moment, which is the safe direction: + // a timeout must never be able to manufacture absence. + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } return nil, err } @@ -191,6 +216,9 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { // The peer answered and completed a handshake but will not carry a // stream, which is a transport condition and never absence. _ = sess.Close() + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } return nil, unreachablePeer(peerID, err) } @@ -199,6 +227,14 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { } // link returns the per-peer entry, creating it on first use. +// +// Entries are never pruned: a peer id opened once keeps its entry, and any +// session cached on it, until Close. The cost is not the map entry. A peer that +// has left the deployment but is still listening keeps a live WebSocket and the +// two yamux loop goroutines behind it for as long as this process runs; a peer +// that is genuinely gone is reclaimed by the 30s keepalive default, so the real +// exposure is narrow. There is no Forget because nothing yet knows which peers +// have left; Task 5's ownership work is where that knowledge appears. func (p *PeerPool) link(peerID string) (*peerLink, error) { p.mu.Lock() defer p.mu.Unlock() @@ -219,6 +255,14 @@ func (p *PeerPool) link(peerID string) (*peerLink, error) { // // A registry miss is returned unchanged so ErrInstanceNotFound reaches the // caller; everything after it is wrapped as unreachable. +// +// The address is only read here, so a peer that re-registers on a new address +// while its current session is still alive keeps being reached over that +// session until it dies. That is deliberate: an address change without a +// session break means the peer is still answering on the old one, and dropping +// a working link to chase a registry write would interrupt live requests for +// nothing. A replica that actually moved breaks its sessions in the process, +// and the re-dial above picks the new address up on the next Open. func (p *PeerPool) dial(ctx context.Context, peerID string) (*yamux.Session, error) { inst, err := p.reg.Get(ctx, peerID) if err != nil { @@ -235,7 +279,7 @@ func (p *PeerPool) dial(ctx context.Context, peerID string) (*yamux.Session, err // link is authenticated by the cluster token rather than by transport. Scheme: "ws", Host: inst.AdvertisedAddr, - Path: clusterep.PeerPath, + Path: PeerPath, RawQuery: url.Values{"id": []string{p.selfID}}.Encode(), } header := http.Header{} @@ -254,7 +298,7 @@ func (p *PeerPool) dial(ctx context.Context, peerID string) (*yamux.Session, err // Client side of the mux: the dialling replica owns the odd stream IDs, // matching the yamux.Server the peer handler puts on its end. - sess, err := yamux.Client(clusterep.WebsocketConn(ws), peerLinkConfig(), nil) + sess, err := yamux.Client(WebsocketConn(ws), peerLinkConfig(), nil) if err != nil { _ = ws.Close() return nil, unreachablePeer(peerID, err) diff --git a/core/services/cluster/peerlink_test.go b/core/services/cluster/peerlink_test.go index 4b145194af06..cfaa4dc69392 100644 --- a/core/services/cluster/peerlink_test.go +++ b/core/services/cluster/peerlink_test.go @@ -3,8 +3,11 @@ package cluster_test import ( "context" "io" + "net" "net/http/httptest" "strings" + "sync" + "time" clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" "github.com/mudler/LocalAI/core/services/cluster" @@ -209,6 +212,77 @@ var _ = Describe("Peer pool", func() { "a reset stream must not cost the whole peer link") }) + It("blames the caller's deadline, not the peer, when a dial runs out of time", func() { + // A listener that completes the TCP connection and then says nothing, + // which is what a peer under load or behind a wedged proxy looks like. + // The peer is not unreachable; the caller is impatient. Reporting + // ErrPeerUnreachable here would make an impatient client enough to get + // a healthy replica routed around. + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = ln.Close() }) + // The accept loop owns the connections it holds and closes them when + // the listener goes away, so nothing is shared with the spec goroutine. + go func() { + defer GinkgoRecover() + var held []net.Conn + defer func() { + for _, c := range held { + _ = c.Close() + } + }() + for { + c, e := ln.Accept() + if e != nil { + return + } + // Hold the connection open without ever answering the upgrade. + held = append(held, c) + } + }() + Expect(reg.Register(ctx, "peer-silent", ln.Addr().String(), "test")).To(Succeed()) + + deadlined, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + DeferCleanup(cancel) + _, err = pool.Open(deadlined, "peer-silent") + Expect(err).To(MatchError(context.DeadlineExceeded)) + Expect(err).ToNot(MatchError(cluster.ErrPeerUnreachable), + "the caller ran out of time; the peer never got a verdict") + Expect(err).ToNot(MatchError(cluster.ErrInstanceNotFound)) + }) + + It("dials once when many callers open the same peer at the same time", func() { + // Without a per-peer lock held across the dial, every concurrent + // caller races to dial and all but one of the resulting sessions is + // dropped on the floor still holding a live WebSocket. + const callers = 16 + streams := make(chan net.Conn, callers) + var wg sync.WaitGroup + for range callers { + wg.Add(1) + go func() { + defer wg.Done() + defer GinkgoRecover() + st, err := pool.Open(context.Background(), "peer-1") + Expect(err).ToNot(HaveOccurred()) + streams <- st + }() + } + wg.Wait() + close(streams) + + count := 0 + for st := range streams { + count++ + DeferCleanup(func(c net.Conn) { _ = c.Close() }, st) + } + Expect(count).To(Equal(callers)) + + Eventually(accepted, "10s").Should(Receive()) + Consistently(accepted, "2s", "200ms").ShouldNot(Receive(), + "concurrent opens must share one dial, not race to dial per caller") + }) + It("refuses to open after Close and is safe to close twice", func() { pool.Close() pool.Close() diff --git a/core/http/endpoints/cluster/wsconn.go b/core/services/cluster/wsconn.go similarity index 100% rename from core/http/endpoints/cluster/wsconn.go rename to core/services/cluster/wsconn.go diff --git a/core/http/endpoints/cluster/wsconn_test.go b/core/services/cluster/wsconn_test.go similarity index 91% rename from core/http/endpoints/cluster/wsconn_test.go rename to core/services/cluster/wsconn_test.go index 5ef3fcf3b1e2..305ce99b3c79 100644 --- a/core/http/endpoints/cluster/wsconn_test.go +++ b/core/services/cluster/wsconn_test.go @@ -12,6 +12,7 @@ import ( "time" clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/gorilla/websocket" "github.com/labstack/echo/v4" @@ -54,8 +55,8 @@ var _ = Describe("WebsocketConn framing", func() { It("returns the rest of a message on the following Read", func() { clientWS, serverWS := wsPair() - writer := clusterep.WebsocketConn(clientWS) - reader := clusterep.WebsocketConn(serverWS) + writer := cluster.WebsocketConn(clientWS) + reader := cluster.WebsocketConn(serverWS) // A lost tail would otherwise park the reassembly below forever; with a // deadline it fails as a timeout on the read that has nothing left. @@ -83,8 +84,8 @@ var _ = Describe("WebsocketConn framing", func() { It("streams a message larger than the yamux read buffer without loss or reordering", func() { clientWS, serverWS := wsPair() - writer := clusterep.WebsocketConn(clientWS) - reader := clusterep.WebsocketConn(serverWS) + writer := cluster.WebsocketConn(clientWS) + reader := cluster.WebsocketConn(serverWS) Expect(reader.SetReadDeadline(time.Now().Add(20 * time.Second))).To(Succeed()) @@ -106,8 +107,8 @@ var _ = Describe("WebsocketConn framing", func() { It("presents consecutive messages as one continuous byte stream", func() { clientWS, serverWS := wsPair() - writer := clusterep.WebsocketConn(clientWS) - reader := clusterep.WebsocketConn(serverWS) + writer := cluster.WebsocketConn(clientWS) + reader := cluster.WebsocketConn(serverWS) Expect(reader.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) @@ -126,8 +127,8 @@ var _ = Describe("WebsocketConn framing", func() { It("never hands back a zero-length read for a zero-length message", func() { clientWS, serverWS := wsPair() - writer := clusterep.WebsocketConn(clientWS) - reader := clusterep.WebsocketConn(serverWS) + writer := cluster.WebsocketConn(clientWS) + reader := cluster.WebsocketConn(serverWS) Expect(reader.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) @@ -146,7 +147,7 @@ var _ = Describe("WebsocketConn framing", func() { It("reports a clean peer close as io.EOF", func() { clientWS, serverWS := wsPair() - reader := clusterep.WebsocketConn(serverWS) + reader := cluster.WebsocketConn(serverWS) Expect(clientWS.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))).To(Succeed()) @@ -157,7 +158,7 @@ var _ = Describe("WebsocketConn framing", func() { It("refuses a text message rather than desynchronising the stream", func() { clientWS, serverWS := wsPair() - reader := clusterep.WebsocketConn(serverWS) + reader := cluster.WebsocketConn(serverWS) Expect(clientWS.WriteMessage(websocket.TextMessage, []byte("not a frame"))).To(Succeed()) @@ -168,7 +169,7 @@ var _ = Describe("WebsocketConn framing", func() { It("satisfies net.Conn", func() { clientWS, _ := wsPair() - var conn net.Conn = clusterep.WebsocketConn(clientWS) + var conn net.Conn = cluster.WebsocketConn(clientWS) Expect(conn.LocalAddr()).ToNot(BeNil()) Expect(conn.RemoteAddr()).ToNot(BeNil()) @@ -176,7 +177,7 @@ var _ = Describe("WebsocketConn framing", func() { It("enforces a write deadline, which yamux arms before every flush", func() { clientWS, _ := wsPair() - conn := clusterep.WebsocketConn(clientWS) + conn := cluster.WebsocketConn(clientWS) // Asserting that the setter returns nil would prove nothing: gorilla // only records the deadline and applies it at the next flush. The write @@ -191,7 +192,7 @@ var _ = Describe("WebsocketConn framing", func() { It("enforces a read deadline, which is how a parked reader is unblocked", func() { clientWS, _ := wsPair() - conn := clusterep.WebsocketConn(clientWS) + conn := cluster.WebsocketConn(clientWS) Expect(conn.SetReadDeadline(time.Now().Add(-time.Second))).To(Succeed()) _, err := conn.Read(make([]byte, 8)) @@ -201,7 +202,7 @@ var _ = Describe("WebsocketConn framing", func() { It("arms both directions from SetDeadline", func() { clientWS, _ := wsPair() - conn := clusterep.WebsocketConn(clientWS) + conn := cluster.WebsocketConn(clientWS) Expect(conn.SetDeadline(time.Now().Add(-time.Second))).To(Succeed()) @@ -217,7 +218,7 @@ var _ = Describe("WebsocketConn framing", func() { // struct field, so this is a data race unless the adapter serialises it; // the spec is here to be run under -race, where it would report one. clientWS, _ := wsPair() - conn := clusterep.WebsocketConn(clientWS) + conn := cluster.WebsocketConn(clientWS) done := make(chan struct{}) go func() { @@ -252,7 +253,7 @@ var _ = Describe("Peer link payloads", func() { var serverSess *yamux.Session Eventually(sessions, "5s").Should(Receive(&serverSess)) - clientSess, err := yamux.Client(clusterep.WebsocketConn(conn), nil, nil) + clientSess, err := yamux.Client(cluster.WebsocketConn(conn), nil, nil) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = clientSess.Close() }) From d9288513e6fd7adc1f92f16ab87b995fd4748406 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 00:50:11 +0000 Subject: [PATCH 11/42] feat(cluster): fence worker-connection ownership with a monotonic epoch A worker whose link is silently broken reconnects to another replica while the old owner's socket has not yet noticed. Without a fence both believe they own it. Claim is a single atomic upsert returning the new epoch, and a release must match both owner and epoch so a stale owner cannot delete a live claim. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/ownership.go | 103 +++++++++++++ core/services/cluster/ownership_test.go | 186 ++++++++++++++++++++++++ core/services/nodes/registry.go | 2 +- 3 files changed, 290 insertions(+), 1 deletion(-) create mode 100644 core/services/cluster/ownership.go create mode 100644 core/services/cluster/ownership_test.go diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go new file mode 100644 index 000000000000..143bc859bbc8 --- /dev/null +++ b/core/services/cluster/ownership.go @@ -0,0 +1,103 @@ +package cluster + +import ( + "context" + "errors" + "fmt" + "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") + +// 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 gets a higher +// epoch than any claim before it, 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. +type NodeConnection struct { + NodeID string `gorm:"primaryKey;size:36" json:"node_id"` + OwnerInstanceID string `gorm:"size:36;index;not null" json:"owner_instance_id"` + Epoch int64 `gorm:"not null" json:"epoch"` + ConnectedAt time.Time `gorm:"not null;default:now()" json:"connected_at"` + LastSeen time.Time `gorm:"index;not null;default:now()" json:"last_seen"` +} + +// Claim records ownerID as the owner of nodeID's tunnel and returns the new +// epoch, which is strictly greater than the epoch of every earlier claim on the +// same node. +// +// 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 increment is +// computed by the database against the row version the winner just wrote. +func (r *Registry) Claim(ctx context.Context, nodeID, ownerID string) (int64, error) { + // Timestamps are stamped by the database, never by this process, for the + // same reason instance liveness is: they are compared across replicas, so + // they have to be measured on the one clock every replica shares. On insert + // that is the column default; on conflict it is the assignment below. + conn := NodeConnection{NodeID: nodeID, OwnerInstanceID: ownerID, Epoch: 1} + if err := r.db.WithContext(ctx).Clauses( + clause.OnConflict{ + Columns: []clause.Column{{Name: "node_id"}}, + DoUpdates: clause.Assignments(map[string]any{ + "owner_instance_id": ownerID, + "epoch": gorm.Expr(`"node_connections"."epoch" + 1`), + "connected_at": gorm.Expr("now()"), + "last_seen": gorm.Expr("now()"), + }), + }, + clause.Returning{Columns: []clause.Column{{Name: "epoch"}}}, + ).Create(&conn).Error; err != nil { + return 0, fmt.Errorf("claiming connection for node %q as %q: %w", nodeID, ownerID, err) + } + return conn.Epoch, nil +} + +// Owner returns the replica that holds nodeID's tunnel and the epoch of that +// claim, or ErrNoConnection when the node has no recorded connection. +func (r *Registry) Owner(ctx context.Context, nodeID string) (string, int64, error) { + var conn NodeConnection + err := r.db.WithContext(ctx).Where("node_id = ?", nodeID).First(&conn).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return "", 0, fmt.Errorf("looking up owner of node %q: %w", nodeID, ErrNoConnection) + } + if err != nil { + return "", 0, fmt.Errorf("looking up owner of node %q: %w", nodeID, err) + } + return conn.OwnerInstanceID, conn.Epoch, nil +} + +// Release drops the claim identified by ownerID and epoch. Both are in the +// WHERE so a replica that has only just noticed its dead socket cannot delete +// the claim a later reconnect established elsewhere: the row it is trying to +// clean up no longer exists, and deleting the live one would strand a worker +// that is in fact connected. 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 { + // gorm reports no error when a Where matches nothing, so the miss has to be + // read off RowsAffected. + res := r.db.WithContext(ctx). + Where("node_id = ? AND owner_instance_id = ? AND epoch = ?", nodeID, ownerID, epoch). + Delete(&NodeConnection{}) + 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 +} diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go new file mode 100644 index 000000000000..cb34b105bff3 --- /dev/null +++ b/core/services/cluster/ownership_test.go @@ -0,0 +1,186 @@ +package cluster_test + +import ( + "context" + "fmt" + "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/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 +} + +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, _ := fc() + r.mu.Lock() + r.statements = append(r.statements, sql) + r.mu.Unlock() +} + +// 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.statements).To(HaveLen(1), "expected exactly one statement, got: %v", r.statements) + return r.statements[0] +} + +var _ = Describe("Connection ownership", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) + reg = cluster.NewRegistry(db) + ctx = context.Background() + }) + + It("increments the epoch on every claim", func() { + 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).To(BeNumerically(">", e1)) + }) + + It("reports the latest owner", func() { + _, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + _, err = reg.Claim(ctx, "w1", "inst-b") + Expect(err).ToNot(HaveOccurred()) + + owner, epoch, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-b")) + Expect(epoch).To(BeNumerically("==", 2)) + }) + + It("distinguishes an unknown connection", func() { + _, _, err := reg.Owner(ctx, "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 := reg.Owner(ctx, "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 := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("inst-a")) + }) + + 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 = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("claims in one statement that increments in SQL 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(ContainSubstring(`"node_connections"."epoch" + 1`), + "the epoch must be incremented by the database, not 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, and its epoch is the highest handed out. + var rows []cluster.NodeConnection + Expect(db.Where("node_id = ?", "w-race").Find(&rows).Error).To(Succeed()) + Expect(rows).To(HaveLen(1)) + var max int64 + for e := range seen { + if e > max { + max = e + } + } + Expect(rows[0].Epoch).To(Equal(max)) + }) +}) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 291ccd2dbb0f..753f26a7b544 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -443,7 +443,7 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}) + return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}, &cluster.NodeConnection{}) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } From 7f1599e83d32d6e615d833aecbc75576dc0aacb2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 01:09:49 +0000 Subject: [PATCH 12/42] fix(cluster): draw connection epochs from a sequence so none is ever reused Release deletes the row, so a per-row `epoch + 1` restarted the numbering at 1 for the next claim. A replica could then be handed an epoch it already held: claim w1 at epoch 1, lose the link silently, watch another replica claim and release, reclaim and be handed 1 again, and its delayed cleanup for the first dead link would match the live claim and delete it. The fence has to be unique per node over time, not per row lifetime. Every claim now draws nextval from a dedicated sequence on both the insert and the conflict paths, so an epoch is never issued twice. The draw still happens after the row lock on the conflict path, so the winning claim still holds the highest epoch handed out. Also drop last_seen. Nothing maintained it and it was always equal to connected_at, but an indexed column named that way invites a second liveness clock; whether the owner is alive is Instance.LastSeen, and whether a claim is current is the epoch. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/ownership.go | 78 +++++++++++++++++++------ core/services/cluster/ownership_test.go | 48 +++++++++++++-- core/services/nodes/registry.go | 9 ++- 3 files changed, 109 insertions(+), 26 deletions(-) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index 143bc859bbc8..ccddd69b9bbb 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -16,54 +16,94 @@ import ( // answer "this worker is not connected here", or to retry. var ErrNoConnection = errors.New("cluster: no connection recorded for node") +// epochSequence is the PostgreSQL sequence every claim draws its epoch from. +// A sequence rather than a per-row counter because a released row is deleted: +// with `epoch = epoch + 1` the numbering restarts 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 delete, 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 gets a higher -// epoch than any claim before it, 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. +// 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"` OwnerInstanceID string `gorm:"size:36;index;not null" json:"owner_instance_id"` Epoch int64 `gorm:"not null" json:"epoch"` ConnectedAt time.Time `gorm:"not null;default:now()" json:"connected_at"` - LastSeen time.Time `gorm:"index;not null;default:now()" json:"last_seen"` +} + +// 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 it 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 { + 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 new -// epoch, which is strictly greater than the epoch of every earlier claim on the -// same node. +// epoch, which is greater than every epoch handed out for that node before it +// and is never handed out again. // // 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 increment is -// computed by the database against the row version the winner just wrote. +// 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) { - // Timestamps are stamped by the database, never by this process, for the - // same reason instance liveness is: they are compared across replicas, so - // they have to be measured on the one clock every replica shares. On insert - // that is the column default; on conflict it is the assignment below. - conn := NodeConnection{NodeID: nodeID, OwnerInstanceID: ownerID, Epoch: 1} - if err := r.db.WithContext(ctx).Clauses( + // 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": gorm.Expr(`"node_connections"."epoch" + 1`), + "epoch": nextEpoch, "connected_at": gorm.Expr("now()"), - "last_seen": gorm.Expr("now()"), }), }, clause.Returning{Columns: []clause.Column{{Name: "epoch"}}}, - ).Create(&conn).Error; err != nil { + ).Create(values).Error; err != nil { return 0, fmt.Errorf("claiming connection for node %q as %q: %w", nodeID, ownerID, err) } - return conn.Epoch, nil + // 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 the epoch of that diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index cb34b105bff3..4e33b056163b 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -24,6 +24,7 @@ type sqlRecorder struct { gormlogger.Interface mu sync.Mutex statements []string + errs []error } func newSQLRecorder() *sqlRecorder { @@ -31,10 +32,17 @@ func newSQLRecorder() *sqlRecorder { } func (r *sqlRecorder) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) { - sql, _ := fc() + 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 @@ -42,6 +50,7 @@ func (r *sqlRecorder) Trace(ctx context.Context, begin time.Time, fc func() (str 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] } @@ -55,9 +64,10 @@ var _ = Describe("Connection ownership", func() { BeforeEach(func() { db = testutil.SetupTestDB() + ctx = context.Background() Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) + Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) reg = cluster.NewRegistry(db) - ctx = context.Background() }) It("increments the epoch on every claim", func() { @@ -71,13 +81,13 @@ var _ = Describe("Connection ownership", func() { It("reports the latest owner", func() { _, err := reg.Claim(ctx, "w1", "inst-a") Expect(err).ToNot(HaveOccurred()) - _, err = reg.Claim(ctx, "w1", "inst-b") + e2, err := reg.Claim(ctx, "w1", "inst-b") Expect(err).ToNot(HaveOccurred()) owner, epoch, err := reg.Owner(ctx, "w1") Expect(err).ToNot(HaveOccurred()) Expect(owner).To(Equal("inst-b")) - Expect(epoch).To(BeNumerically("==", 2)) + Expect(epoch).To(Equal(e2), "the stored epoch must be the one the winning claim was handed") }) It("distinguishes an unknown connection", func() { @@ -114,6 +124,32 @@ var _ = Describe("Connection ownership", func() { 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 := reg.Owner(ctx, "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()) @@ -135,8 +171,8 @@ var _ = Describe("Connection ownership", func() { // 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(ContainSubstring(`"node_connections"."epoch" + 1`), - "the epoch must be incremented by the database, not by this process") + 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 diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 753f26a7b544..72b957bcc27a 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -443,7 +443,14 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}, &cluster.NodeConnection{}) + if err := db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}, &cluster.NodeConnection{}); err != nil { + return err + } + // AutoMigrate models tables and columns but has no notion of a + // sequence, and the connection-ownership fence draws its epochs from + // one. It runs under this same lock so concurrently starting replicas + // do not race on the DDL. + return cluster.EnsureEpochSequence(context.Background(), db) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } From 667ce1d9a9772234ae542ed46288560c30304ff6 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 01:28:27 +0000 Subject: [PATCH 13/42] fix(cluster): keep the connection schema migratable on SQLite The NodeConnection model carried `default:now()`, which is PostgreSQL syntax reaching the DDL, so AutoMigrate failed on the single-binary SQLite path and took every SQLite caller of nodes.NewNodeRegistry down with it. Stamp the database clock as an expression inside Claim instead, the way Register already does, and leave the column plain. CREATE SEQUENCE is Postgres-only for the same reason, so it is skipped on another dialect, and Claim refuses that dialect outright: a fence that cannot draw a token must say so rather than fail later as a missing function. Also correct a claim the previous commit made in both the doc comment and its message. An epoch is unique and never reissued, but it is not ordered: the insert path draws its sequence value before taking the row lock, so a claim that inserts after a release can be handed a lower number than one already issued. Uniqueness is what Release needs, since it matches by equality; callers must never compare epochs for order. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/ownership.go | 52 +++++++++++++++++++++---- core/services/cluster/ownership_test.go | 38 +++++++++++++++++- 2 files changed, 82 insertions(+), 8 deletions(-) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index ccddd69b9bbb..a2e2ba9cb5ae 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "strings" "time" "gorm.io/gorm" @@ -16,6 +17,16 @@ import ( // 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 released row is deleted: // with `epoch = epoch + 1` the numbering restarts at 1 for the next claim, so a @@ -39,10 +50,13 @@ const epochSequence = "node_connection_epochs" // 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"` - OwnerInstanceID string `gorm:"size:36;index;not null" json:"owner_instance_id"` - Epoch int64 `gorm:"not null" json:"epoch"` - ConnectedAt time.Time `gorm:"not null;default:now()" json:"connected_at"` + NodeID string `gorm:"primaryKey;size:36" json:"node_id"` + 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"` } // EnsureEpochSequence creates the sequence Claim draws epochs from. It lives @@ -56,15 +70,32 @@ type NodeConnection struct { // (`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 new -// epoch, which is greater than every epoch handed out for that node before it -// and is never handed out again. +// 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 @@ -73,6 +104,13 @@ func EnsureEpochSequence(ctx context.Context, db *gorm.DB) error { // 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. diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index 4e33b056163b..ba634ebcd3eb 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -3,6 +3,7 @@ package cluster_test import ( "context" "fmt" + "path/filepath" "strings" "sync" "time" @@ -12,6 +13,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "gorm.io/driver/sqlite" "gorm.io/gorm" gormlogger "gorm.io/gorm/logger" ) @@ -207,7 +209,11 @@ var _ = Describe("Connection ownership", func() { } Expect(seen).To(HaveLen(claimants)) - // Exactly one row, and its epoch is the highest handed out. + // Exactly one row, and its epoch is the highest handed out. That last + // part holds here because no Release intervenes: every claim after the + // first blocks on the row lock and draws its epoch after taking it, in + // commit order. It is not a general guarantee about epochs, which are + // unique but unordered. var rows []cluster.NodeConnection Expect(db.Where("node_id = ?", "w-race").Find(&rows).Error).To(Succeed()) Expect(rows).To(HaveLen(1)) @@ -220,3 +226,33 @@ var _ = Describe("Connection ownership", func() { Expect(rows[0].Epoch).To(Equal(max)) }) }) + +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(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) + Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + }) + + It("refuses to claim, rather than pretending to fence", func() { + Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) + Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + + _, err := cluster.NewRegistry(db).Claim(ctx, "w1", "inst-a") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("requires PostgreSQL")) + }) +}) From aca383d263838cd29f3a7e916b35cfe38490ca10 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 02:28:42 +0000 Subject: [PATCH 14/42] feat(cluster): give phase 1 a call site, and prove it against real replicas Tasks 1 to 5 built an instances table, a splice, both halves of a peer link and an epoch fence, and nothing in the tree called any of it: no replica registered, no route was mounted, no sweeper ran. Proving phase 1 end to end therefore had to start by wiring it. A frontend in distributed mode now publishes the address its peers dial, heartbeats it, and sweeps replicas that stopped answering along with the connection rows they owned, in one pass so the two can never disagree about who is alive. It serves the peer link and owns the sessions peers dial in, refusing streams on them until phase 2 installs a relay: a session nobody accepts on does not fail a peer's Open, it hangs it. The address is the one peers use, not the one the process binds, and it is derived from the route to PostgreSQL. That derivation only holds while the database is remote, so LOCALAI_DISTRIBUTED_ADVERTISE_ADDR sets it explicitly and a replica that can determine neither warns and keeps serving rather than failing to start. Three e2e scenarios run against real local-ai processes, real PostgreSQL and real dials: replicas publish addresses that can actually be connected to; a sibling opens a stream over the peer link and is refused without the cluster token; and a killed replica is reported unreachable, never absent, loses the claim it held, and takes no worker with it. Each was verified by mutation: eight injected defects, each failing the scenario that claims to catch it. Also moves RegisterClusterRoutes to core/http/routes beside every other registrar, folds AutoMigrate and the epoch sequence into one cluster.Migrate, and turns the peer route's auth-coverage spec into a real assertion: it drives the request through the actual auth middleware instead of comparing two string constants, which the old spec would have passed even with the exemption deleted. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 73 +++++ core/cli/run.go | 4 + core/config/distributed_config.go | 19 +- core/http/app.go | 8 + core/http/endpoints/cluster/peer.go | 12 - core/http/endpoints/cluster/peer_test.go | 110 ++++++-- core/http/routes/cluster.go | 25 ++ core/services/cluster/instance_test.go | 4 +- core/services/cluster/membership.go | 175 ++++++++++++ core/services/cluster/membership_test.go | 116 ++++++++ core/services/cluster/ownership.go | 25 +- core/services/cluster/ownership_test.go | 9 +- core/services/cluster/peerlink_test.go | 18 +- core/services/cluster/sessions.go | 140 ++++++++++ core/services/cluster/sessions_test.go | 122 ++++++++ core/services/cluster/wsconn_test.go | 3 +- core/services/nodes/registry.go | 12 +- docs/content/features/distributed-mode.md | 22 ++ tests/e2e/distributed/cluster/cluster.go | 16 ++ .../e2e/distributed/cluster_baseline_test.go | 10 +- .../e2e/distributed/cluster_peerlink_test.go | 263 ++++++++++++++++++ 21 files changed, 1119 insertions(+), 67 deletions(-) create mode 100644 core/http/routes/cluster.go create mode 100644 core/services/cluster/membership.go create mode 100644 core/services/cluster/membership_test.go create mode 100644 core/services/cluster/sessions.go create mode 100644 core/services/cluster/sessions_test.go create mode 100644 tests/e2e/distributed/cluster_peerlink_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index b7dc0bf91351..404b2112c3fc 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "io" + "net" + "strconv" "strings" "sync" "time" @@ -12,12 +14,14 @@ import ( "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/agents" + "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" "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" @@ -43,6 +47,16 @@ type DistributedServices struct { Unloader *nodes.RemoteUnloaderAdapter ModelCleanup *nodes.ModelCleanupService + // Cluster is the replica-membership registry: which frontend replicas are + // alive, at which address, and which of them holds a given worker's tunnel. + Cluster *cluster.Registry + // Membership publishes this replica's row and reaps the dead. Nil when no + // peer-reachable address could be determined, which leaves this replica + // invisible to its peers but otherwise fully functional. + Membership *cluster.Membership + // PeerSessions owns the peer links other replicas dialled into this one. + PeerSessions *cluster.SessionStore + shutdownOnce sync.Once } @@ -53,6 +67,15 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { + // Peer state first: a replica that is going away should stop claiming + // to be alive before it stops answering, so peers re-home rather than + // dial a process in teardown. + if ds.Membership != nil { + ds.Membership.Stop() + } + if ds.PeerSessions != nil { + ds.PeerSessions.CloseAll() + } if ds.Health != nil { ds.Health.Stop() } @@ -162,6 +185,29 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Node registry initialized") + // Replica membership. NewNodeRegistry has just migrated the tables this + // reads, so it has to come after it. + clusterRegistry := cluster.NewRegistry(authDB) + // Accepted peer links are held with no stream handler: this replica has + // somewhere to put a link a peer dials, and refuses the streams on it, + // because nothing relays worker traffic yet. + peerSessions := cluster.NewSessionStore(nil) + var membership *cluster.Membership + if advertised, err := advertisedPeerAddr(cfg); err != nil { + // Not fatal. A replica that cannot publish an address still serves + // every request that reaches it directly; what it cannot do is have + // another replica relay to it. Failing startup here would take out + // every existing single-host deployment, whose route to a local + // database is loopback. + xlog.Warn("This replica will not be reachable by its peers: no advertised address", + "error", err, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + } else { + membership = cluster.NewMembership(clusterRegistry, cfg.Distributed.InstanceID, advertised, internal.PrintableVersion()) + if err := membership.Start(cfg.Context); err != nil { + return nil, fmt.Errorf("registering this replica in the cluster: %w", err) + } + } + // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an // operator can pin placement to a stable name like "production" and have it @@ -450,9 +496,36 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade ModelAdapter: modelAdapter, Unloader: remoteUnloader, ModelCleanup: modelCleanup, + Cluster: clusterRegistry, + Membership: membership, + PeerSessions: peerSessions, }, nil } +// advertisedPeerAddr is the host:port peers dial to reach this replica. +// +// The operator's value wins outright. Otherwise it is derived from the port +// this process serves on and the local address that routes to PostgreSQL, which +// is only a peer-reachable answer when the database is on another host; +// DiscoverAdvertisedAddr refuses rather than guessing when it is not. +func advertisedPeerAddr(cfg *config.ApplicationConfig) (string, error) { + if cfg.Distributed.AdvertiseAddr != "" { + return cfg.Distributed.AdvertiseAddr, nil + } + if cfg.APIAddress == "" { + return "", fmt.Errorf("no API address to derive a peer port from") + } + _, port, err := net.SplitHostPort(cfg.APIAddress) + if err != nil { + return "", fmt.Errorf("reading the peer port out of API address %q: %w", cfg.APIAddress, err) + } + portNumber, err := strconv.Atoi(port) + if err != nil { + return "", fmt.Errorf("API address %q has a non-numeric port: %w", cfg.APIAddress, err) + } + return cluster.DiscoverAdvertisedAddr(cfg.Auth.DatabaseURL, portNumber) +} + func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } diff --git a/core/cli/run.go b/core/cli/run.go index 6b9b3e4dc0b9..4408edd1203c 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -165,6 +165,7 @@ 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"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` + DistributedAdvertiseAddr string `env:"LOCALAI_DISTRIBUTED_ADVERTISE_ADDR" help:"host:port other frontend replicas dial to reach this one (peer link). Empty = derived from the local address that routes to PostgreSQL, which only works when the database is on another host." 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"` @@ -351,6 +352,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.InstanceID != "" { opts = append(opts, config.WithDistributedInstanceID(r.InstanceID)) } + if r.DistributedAdvertiseAddr != "" { + opts = append(opts, config.WithDistributedAdvertiseAddr(r.DistributedAdvertiseAddr)) + } 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 5a48a84e9b44..bbb141592934 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -13,8 +13,15 @@ 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) + // AdvertiseAddr is the host:port OTHER REPLICAS dial to reach this one, + // which is not the address this process binds: a replica behind a service + // or a NAT binds one and is reached at another. Empty means "work it out", + // by asking the kernel which local address routes to PostgreSQL; that + // answer is only usable when the database is remote, so a deployment with + // a local or sidecar database has to set this. + AdvertiseAddr string // LOCALAI_DISTRIBUTED_ADVERTISE_ADDR 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) @@ -195,6 +202,14 @@ func WithDistributedInstanceID(id string) AppOption { } } +// WithDistributedAdvertiseAddr pins the host:port peers dial to reach this +// replica, overriding the route-based discovery. +func WithDistributedAdvertiseAddr(addr string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.AdvertiseAddr = 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 2e1453ac0e38..d6418165b5f0 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -576,6 +576,14 @@ func API(application *application.Application) (*echo.Echo, error) { 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) + // Replica-to-replica peer link. Registered only in distributed mode: in + // single-node mode there are no peers, and the route authenticates with the + // registration token, so publishing it unconditionally would put a + // multiplexer on every single-binary install. + if d := application.Distributed(); d != nil && d.PeerSessions != nil { + routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, d.PeerSessions.Accept) + } + // Distributed SSE routes (job progress + agent events via NATS) if d := application.Distributed(); d != nil { if d.Dispatcher != nil { diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go index 85985eee5762..f0fb96ebc16e 100644 --- a/core/http/endpoints/cluster/peer.go +++ b/core/http/endpoints/cluster/peer.go @@ -18,18 +18,6 @@ import ( "github.com/mudler/xlog" ) -// RegisterClusterRoutes registers the peer link. onPeer receives every -// authenticated session; see PeerHandler for what it is expected to do with it. -// -// The route is core/services/cluster's own constant, so the handler and the -// dialler cannot be registered and dialled at different paths. That the path -// also falls under auth.ClusterPathPrefix, and so bypasses the session -// middleware, is asserted by a spec in this package: it is the only place that -// can see both, since core/services/cluster must not import core/http/auth. -func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { - e.GET(clustersvc.PeerPath, PeerHandler(token, onPeer)) -} - // PeerHandler upgrades an authenticated peer dial to a WebSocket, wraps it as // a yamux server session and hands it to onSession. // diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go index a86ab8add8bd..dc8cba3da9e6 100644 --- a/core/http/endpoints/cluster/peer_test.go +++ b/core/http/endpoints/cluster/peer_test.go @@ -5,8 +5,9 @@ import ( "net/http/httptest" "strings" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/auth" - clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/http/routes" clustersvc "github.com/mudler/LocalAI/core/services/cluster" "github.com/gorilla/websocket" @@ -16,6 +17,11 @@ import ( . "github.com/onsi/gomega" ) +// wsPeerURL is the peer route on a test server, named as peer-1. +func wsPeerURL(s *httptest.Server) string { + return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.PeerPath + "?id=peer-1" +} + var _ = Describe("Peer link handler", func() { var ( srv *httptest.Server @@ -25,19 +31,15 @@ var _ = Describe("Peer link handler", func() { BeforeEach(func() { sessions = make(chan *yamux.Session, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) srv = httptest.NewServer(e) DeferCleanup(srv.Close) }) - wsURL := func(s *httptest.Server) string { - return "ws" + strings.TrimPrefix(s.URL, "http") + "/api/cluster/peer?id=peer-1" - } - It("rejects a connection with no token", func() { - _, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), nil) + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), nil) Expect(err).To(HaveOccurred()) Expect(resp).ToNot(BeNil()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) @@ -46,7 +48,7 @@ var _ = Describe("Peer link handler", func() { It("rejects a connection with the wrong token", func() { h := http.Header{} h.Set("Authorization", "Bearer wrong") - _, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) Expect(err).To(HaveOccurred()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) }) @@ -54,7 +56,7 @@ var _ = Describe("Peer link handler", func() { It("accepts an authenticated peer and yields a usable yamux session", func() { h := http.Header{} h.Set("Authorization", "Bearer peer-token") - conn, _, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) @@ -98,11 +100,11 @@ var _ = Describe("Peer link handler", func() { h.Set("Authorization", "Bearer peer-token") ids := make(chan string, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + routes.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) s2 := httptest.NewServer(e) DeferCleanup(s2.Close) - conn, _, err := websocket.DefaultDialer.Dial(wsURL(s2), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) @@ -115,12 +117,12 @@ var _ = Describe("Peer link handler", func() { // unauthenticated yamux multiplexer to anyone who can reach the port. e := echo.New() accepted := make(chan *yamux.Session, 1) - clusterep.RegisterClusterRoutes(e, "", func(_ string, sess *yamux.Session) { accepted <- sess }) + routes.RegisterClusterRoutes(e, "", func(_ string, sess *yamux.Session) { accepted <- sess }) s2 := httptest.NewServer(e) DeferCleanup(s2.Close) for _, header := range []http.Header{nil, {"Authorization": []string{"Bearer "}}, {"Authorization": []string{"Bearer anything"}}} { - _, resp, err := websocket.DefaultDialer.Dial(wsURL(s2), header) + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), header) Expect(err).To(HaveOccurred()) Expect(resp).ToNot(BeNil()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) @@ -132,7 +134,7 @@ var _ = Describe("Peer link handler", func() { // RFC 7235 makes the scheme case-insensitive. The token after it is not. h := http.Header{} h.Set("Authorization", "bearer peer-token") - conn, _, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) Eventually(sessions, "5s").Should(Receive()) @@ -143,7 +145,7 @@ var _ = Describe("Peer link handler", func() { // without the handler's own recover the peer would keep a link nobody // ever accepts streams on. e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, _ *yamux.Session) { + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, _ *yamux.Session) { panic("callback exploded") }) s2 := httptest.NewServer(e) @@ -157,7 +159,7 @@ var _ = Describe("Peer link handler", func() { h := http.Header{} h.Set("Authorization", "Bearer peer-token") - conn, _, err := websocket.DefaultDialer.Dial(wsURL(s2), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) @@ -186,17 +188,69 @@ var _ = Describe("Peer link handler", func() { }) }) -var _ = Describe("Peer link auth prefix", func() { - It("keeps the peer route inside the alternative-authentication prefix", func() { - // The peer route authenticates with the cluster token, not the global - // session middleware, which only holds while the route sits under the - // prefix auth exempts. Moving either one alone 401s every peer dial. - // - // This lives here because it is the only package that can see both: - // core/services/cluster owns the route and must stay free of any - // core/http dependency, and core/http/auth owns the exemption. - Expect(strings.HasPrefix(clustersvc.PeerPath, auth.ClusterPathPrefix)).To(BeTrue(), - "peer route %q is no longer under the auth-exempt prefix %q", - clustersvc.PeerPath, auth.ClusterPathPrefix) +var _ = Describe("Peer link auth coverage", func() { + // These specs put the REAL global auth middleware in front of the REAL + // registrar and prove a peer dial reaches the handler anyway. The peer link + // authenticates with the cluster token, not a session, so it only works + // while its path sits under the prefix auth exempts; moving either one + // alone 401s every peer dial, and the two live in packages that must not + // import each other. + // + // The predicate that grants the exemption is unexported, so this asserts on + // its effect rather than on it: what a caller can observe is whether the + // request reaches the handler. + var ( + srv *httptest.Server + sessions chan *yamux.Session + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + e := echo.New() + // A nil DB with one legacy API key is the cheapest configuration that + // turns the middleware ON without a database. With neither, Middleware + // short-circuits to next() and every assertion below would pass against + // a server that has no auth at all. + e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}})) + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) + // A route outside the cluster prefix, registered on the same server, is + // the control: it proves the middleware in front of both is live. + e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an uncredentialed request to a route outside the cluster prefix", func() { + resp, err := http.Get(srv.URL + "/api/nodes") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), + "the global auth middleware is not actually guarding this server, so the peer-route assertions below would prove nothing") + }) + + It("lets a peer dial reach the handler, which is the only thing that can authenticate it", func() { + // The cluster token is not one of the API keys the middleware knows, so + // a 400 from the handler's own missing-id check can only mean the + // request was let through unauthenticated by the middleware. + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.PeerPath, nil) + Expect(err).ToNot(HaveOccurred()) + req.Header.Set("Authorization", "Bearer peer-token") + + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest), + "a peer dial must reach the handler; 401 here means the peer route left the auth-exempt prefix %q", auth.ClusterPathPrefix) + }) + + It("completes a full peer handshake through the guarded server", func() { + // The status-code assertion above cannot see the upgrade, and the + // upgrade is what a peer actually does. + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Eventually(sessions, "5s").Should(Receive()) }) }) diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go new file mode 100644 index 000000000000..3e462db1cd75 --- /dev/null +++ b/core/http/routes/cluster.go @@ -0,0 +1,25 @@ +package routes + +import ( + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" +) + +// RegisterClusterRoutes registers the replica-to-replica peer link. onPeer +// receives every authenticated session; see clusterep.PeerHandler for what it +// is expected to do with it. +// +// The path is core/services/cluster's own constant, so the handler and the +// dialler cannot be registered and dialled at different paths. That the path +// also falls under auth.ClusterPathPrefix, and so bypasses the global session +// middleware, is asserted by driving a request through that middleware in +// core/http/endpoints/cluster/peer_test.go. +// +// The route carries no auth middleware: it authenticates itself against the +// cluster token, because a peer replica has no session and no user. +func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { + e.GET(clustersvc.PeerPath, clusterep.PeerHandler(token, onPeer)) +} diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go index 8e74e930e3b9..d8075a7a3cad 100644 --- a/core/services/cluster/instance_test.go +++ b/core/services/cluster/instance_test.go @@ -22,9 +22,9 @@ var _ = Describe("Instance registry", func() { BeforeEach(func() { db = testutil.SetupTestDB() - Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) - reg = cluster.NewRegistry(db) ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) }) It("registers an instance and reads it back", func() { diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go new file mode 100644 index 000000000000..3b1182b17458 --- /dev/null +++ b/core/services/cluster/membership.go @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "context" + "errors" + "fmt" + "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 consecutive misses. + // + // The window is generous on purpose. Declaring a replica dead deletes the + // connection rows it owned, and a worker whose row is deleted while its + // owner is merely slow has to be re-homed for nothing. The cost of waiting + // is bounded and symmetric: traffic for that worker is retried, not lost. + InstanceLiveness = 30 * time.Second +) + +// Membership publishes this replica's address and keeps the instances table +// free of replicas that have stopped answering. +// +// It is the only writer of this replica's row and the only sweeper of anyone +// else's, which is what 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 + addr string + version string + + interval time.Duration + liveness time.Duration + + stop chan struct{} + done chan struct{} + stopOnce sync.Once +} + +// NewMembership returns the membership loop for one replica. The address is +// what peers will dial, so it must be reachable from another host, not the +// address this process binds. +func NewMembership(reg *Registry, id, addr, version string) *Membership { + return &Membership{ + reg: reg, + id: id, + addr: addr, + version: version, + interval: InstanceHeartbeat, + liveness: InstanceLiveness, + stop: make(chan struct{}), + done: make(chan struct{}), + } +} + +// Start registers this replica and begins heartbeating and sweeping. The first +// registration is synchronous and its failure is returned: a replica whose +// address never reaches the table is invisible to its peers, and starting +// anyway would hide that behind a background log line. +func (m *Membership) Start(ctx context.Context) error { + if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { + return err + } + xlog.Info("Cluster instance registered", "id", m.id, "addr", m.addr) + go m.loop(ctx) + return nil +} + +// Stop ends the loop and waits for it. Safe to call more than once. +func (m *Membership) Stop() { + if m == nil { + return + } + m.stopOnce.Do(func() { close(m.stop) }) + <-m.done +} + +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 this replica's row and sweeps the dead. +// +// Every replica sweeps, rather than one elected sweeper. The deletes are +// idempotent and cheap, and an elected sweeper is one more thing that has to be +// alive for the cluster to notice that something 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 was stalled long + // enough to look dead. Re-register rather than heartbeat: a heartbeat + // carries no address, so the row has to be rebuilt from scratch. + xlog.Warn("Cluster instance row was reaped, re-registering", "id", m.id) + if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { + xlog.Error("Re-registering cluster instance failed", "id", m.id, "error", err) + } + } else if err != nil { + xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) + } + + instances, connections, 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 || connections > 0 { + xlog.Info("Reaped cluster state left by dead replicas", "instances", instances, "connections", connections) + } +} + +// ReapStale deletes the replicas that have not heartbeated within the liveness +// window, and the connection rows whose owner is no longer among the survivors. +// +// The two deletes are one sweeper on purpose. A connection row is only ever +// orphaned by its owner dying, so the moment that is decided is the moment to +// clean up after it; a second sweeper with its own schedule would either lag +// this one or race it, and would need its own answer to "is that replica +// alive", which is the one fact this table already owns. +// +// self is never reaped. This process may fail to heartbeat for longer than the +// window (a long stall, a database blip) and still be serving: deleting its own +// row would then delete the connections of workers that are, at that moment, +// connected to it. The stall is recovered by the re-register in tick instead. +// +// PostgreSQL only, like Live: distributed mode requires it, and the interval +// arithmetic is measured on the database's clock because liveness is compared +// across replicas. +func (r *Registry) ReapStale(ctx context.Context, self string, within time.Duration) (instances int64, connections int64, err error) { + err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + res := tx.Where("id <> ? AND last_seen <= now() - make_interval(secs => ?)", self, within.Seconds()). + Delete(&Instance{}) + if res.Error != nil { + return fmt.Errorf("deleting stale instances: %w", res.Error) + } + instances = res.RowsAffected + + // Whatever survived the delete above is the live set, so this needs no + // second liveness rule and cannot disagree with the first one. + res = tx.Where("owner_instance_id NOT IN (SELECT id FROM instances)"). + Delete(&NodeConnection{}) + if res.Error != nil { + return fmt.Errorf("deleting orphaned node connections: %w", res.Error) + } + connections = res.RowsAffected + return nil + }) + if err != nil { + return 0, 0, fmt.Errorf("reaping stale cluster state: %w", err) + } + return instances, connections, nil +} diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go new file mode 100644 index 000000000000..a00a08c0b954 --- /dev/null +++ b/core/services/cluster/membership_test.go @@ -0,0 +1,116 @@ +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", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "dead", "10.0.0.2:8080", "v1")).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 = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("leaves the connections of a live replica alone", func() { + Expect(reg.Register(ctx, "live", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "other", "10.0.0.2:8080", "v1")).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 := reg.Owner(ctx, "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", "10.0.0.1:8080", "v1")).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 := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me")) + }) + + It("keeps this replica's row alive and reaps the dead while it runs", func() { + Expect(reg.Register(ctx, "dead", "10.0.0.2:8080", "v1")).To(Succeed()) + age("dead", time.Hour) + + membership := cluster.NewMembership(reg, "me", "10.0.0.1:8080", "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") + }) +}) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index a2e2ba9cb5ae..3b6b6cb622e4 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -59,17 +59,34 @@ type NodeConnection struct { ConnectedAt time.Time `gorm:"not null" json:"connected_at"` } -// EnsureEpochSequence creates the sequence Claim draws epochs from. It lives +// 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 it so that concurrently starting replicas do -// not race on the DDL. It is safe to call repeatedly. +// 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 { +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 diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index ba634ebcd3eb..5744e07cd170 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -67,8 +67,7 @@ var _ = Describe("Connection ownership", func() { BeforeEach(func() { db = testutil.SetupTestDB() ctx = context.Background() - Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) - Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) reg = cluster.NewRegistry(db) }) @@ -243,13 +242,11 @@ var _ = Describe("Connection ownership on a non-PostgreSQL dialect", func() { 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(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) - Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) }) It("refuses to claim, rather than pretending to fence", func() { - Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) - Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) _, err := cluster.NewRegistry(db).Claim(ctx, "w1", "inst-a") Expect(err).To(HaveOccurred()) diff --git a/core/services/cluster/peerlink_test.go b/core/services/cluster/peerlink_test.go index cfaa4dc69392..692f2ee6c22a 100644 --- a/core/services/cluster/peerlink_test.go +++ b/core/services/cluster/peerlink_test.go @@ -20,6 +20,16 @@ import ( "gorm.io/gorm" ) +// servePeerRoute mounts the peer handler on the route both sides agree on. +// +// It deliberately does not call routes.RegisterClusterRoutes: that registrar +// lives in core/http/routes, which imports half the server, and these specs are +// about the handler and the dialler rather than about the route table. The path +// comes from the same constant the registrar uses, so the two cannot drift. +func servePeerRoute(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { + e.GET(cluster.PeerPath, clusterep.PeerHandler(token, onPeer)) +} + var _ = Describe("Peer pool", func() { var ( db *gorm.DB @@ -33,7 +43,7 @@ var _ = Describe("Peer pool", func() { // startPeer stands up a real peer server and registers it under peerID. startPeer := func(peerID string) *httptest.Server { e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + servePeerRoute(e, "peer-token", func(_ string, s *yamux.Session) { accepted <- s }) ts := httptest.NewServer(e) @@ -45,7 +55,7 @@ var _ = Describe("Peer pool", func() { BeforeEach(func() { ctx = context.Background() db = testutil.SetupTestDB() - Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) reg = cluster.NewRegistry(db) accepted = make(chan *yamux.Session, 4) pool = cluster.NewPeerPool("self", "peer-token", reg) @@ -88,7 +98,7 @@ var _ = Describe("Peer pool", func() { // link anonymous and indistinguishable from every other. ids := make(chan string, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + servePeerRoute(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) ts := httptest.NewServer(e) DeferCleanup(ts.Close) Expect(reg.Register(ctx, "peer-named", strings.TrimPrefix(ts.URL, "http://"), "test")).To(Succeed()) @@ -131,7 +141,7 @@ var _ = Describe("Peer pool", func() { // row. Reporting absence here would evict every worker behind a peer // that was merely rolled out with a stale secret. e := echo.New() - clusterep.RegisterClusterRoutes(e, "a-different-token", func(_ string, s *yamux.Session) { accepted <- s }) + servePeerRoute(e, "a-different-token", func(_ string, s *yamux.Session) { accepted <- s }) ts := httptest.NewServer(e) DeferCleanup(ts.Close) Expect(reg.Register(ctx, "peer-strict", strings.TrimPrefix(ts.URL, "http://"), "test")).To(Succeed()) diff --git a/core/services/cluster/sessions.go b/core/services/cluster/sessions.go new file mode 100644 index 000000000000..7ee5fa3327af --- /dev/null +++ b/core/services/cluster/sessions.go @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "net" + "sync" + + "github.com/libp2p/go-yamux/v5" + "github.com/mudler/xlog" +) + +// SessionStore holds the peer links this replica has ACCEPTED, which is the +// mirror image of PeerPool: the pool owns the sessions this replica dialled, +// this owns the ones its peers dialled into it. +// +// Something has to own an accepted session. The HTTP handler cannot: it returns +// as soon as the upgrade is done, and the hijacked connection outlives it. And +// something has to accept the streams that arrive on it, because yamux only +// acknowledges a stream once the far side accepts it, so a session nobody +// accepts on does not fail a peer's Open, it hangs it. +type SessionStore struct { + // onStream handles one accepted stream and owns closing it. A nil handler + // closes the stream immediately, which is what a replica with no relay + // installed should do: refuse promptly rather than leave a peer parked. + onStream func(peerID string, stream net.Conn) + + mu sync.Mutex + sessions map[string]*yamux.Session + closed bool +} + +// NewSessionStore returns a store whose accepted streams are handled by +// onStream. Pass nil to refuse every stream, closing it at once. +func NewSessionStore(onStream func(peerID string, stream net.Conn)) *SessionStore { + return &SessionStore{onStream: onStream, sessions: map[string]*yamux.Session{}} +} + +// Accept takes ownership of a session a peer dialled in. It is the callback +// shape RegisterClusterRoutes wants, and it returns promptly: the serving loop +// runs on its own goroutine, because the handler's return is what completes the +// hijack. +func (s *SessionStore) Accept(peerID string, sess *yamux.Session) { + if sess == nil { + return + } + + s.mu.Lock() + if s.closed { + s.mu.Unlock() + // Shutdown raced the dial. Leaving the session open would keep the peer + // believing it has a live link into a process that is going away. + _ = 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 has noticed. Keeping both would leave a session nothing can ever be + // routed to, since the map holds one per peer. + if previous != nil { + xlog.Debug("cluster peer re-dialled, dropping its previous link", "peer", peerID) + _ = previous.Close() + } + + go s.serve(peerID, sess) +} + +// Get returns the session this replica accepted from peerID. The second result +// is false when no link from that peer is held, which a caller must not read as +// the peer being absent: it may be about to dial, or dialling this replica may +// simply not be its job. +func (s *SessionStore) Get(peerID string) (*yamux.Session, bool) { + s.mu.Lock() + defer s.mu.Unlock() + sess, ok := s.sessions[peerID] + return sess, ok +} + +// serve accepts streams until the session dies, then forgets it. +func (s *SessionStore) serve(peerID string, sess *yamux.Session) { + defer func() { + s.forget(peerID, sess) + _ = sess.Close() + }() + + for { + stream, err := sess.AcceptStream() + if err != nil { + // A peer link ending is ordinary: a rolling update closes every + // session it holds. The error is the session's, not one stream's, + // so there is nothing to recover to. + xlog.Debug("cluster peer link ended", "peer", peerID, "error", err) + return + } + if s.onStream == nil { + // No relay installed. Closing is deliberate and is not the same as + // ignoring: a stream nobody answers parks the peer's request until + // its own deadline, and reports nothing about why. + xlog.Debug("cluster peer stream refused: no relay installed", "peer", peerID) + _ = stream.Close() + continue + } + // One goroutine per stream: the handler relays a whole request, and + // serving them from the accept loop would let one request stall every + // other stream on the link. + go s.onStream(peerID, stream) + } +} + +// forget drops the entry only if it still names this session. A peer that +// re-dialled has already replaced it, and deleting blindly would evict the live +// link when the old one finally noticed it was dead. +func (s *SessionStore) forget(peerID string, sess *yamux.Session) { + s.mu.Lock() + defer s.mu.Unlock() + if s.sessions[peerID] == sess { + delete(s.sessions, peerID) + } +} + +// CloseAll drops every held session. An Accept after it closes the session +// rather than storing it, so a dial racing shutdown cannot leak a link. +func (s *SessionStore) CloseAll() { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return + } + s.closed = true + held := s.sessions + s.sessions = map[string]*yamux.Session{} + s.mu.Unlock() + + for _, sess := range held { + _ = sess.Close() + } +} diff --git a/core/services/cluster/sessions_test.go b/core/services/cluster/sessions_test.go new file mode 100644 index 000000000000..f3fb44ca8cae --- /dev/null +++ b/core/services/cluster/sessions_test.go @@ -0,0 +1,122 @@ +package cluster_test + +import ( + "net" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// yamuxPair returns a client and a server session over an in-memory pipe. It +// stands in for a dialled peer link: everything the store does with a session +// is transport-agnostic, and the WebSocket half is covered where it is used. +func yamuxPair() (client *yamux.Session, server *yamux.Session) { + GinkgoHelper() + a, b := net.Pipe() + var err error + server, err = yamux.Server(a, nil, nil) + Expect(err).ToNot(HaveOccurred()) + client, err = yamux.Client(b, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { + _ = client.Close() + _ = server.Close() + }) + return client, server +} + +var _ = Describe("Accepted peer sessions", func() { + It("accepts and refuses a stream rather than leaving the peer parked", func() { + // yamux only acknowledges a stream once the far side accepts it, so a + // store that held the session without accepting would not fail a peer's + // Open, it would hang it, and every relayed request behind it. + store := cluster.NewSessionStore(nil) + DeferCleanup(store.CloseAll) + client, server := yamuxPair() + store.Accept("peer-1", server) + + stream, err := client.OpenStream(GinkgoT().Context()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + Expect(stream.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = stream.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred(), "a refused stream must end, not hang") + }) + + It("hands a stream to the relay when one is installed", func() { + streams := make(chan net.Conn, 1) + store := cluster.NewSessionStore(func(_ string, stream net.Conn) { streams <- stream }) + DeferCleanup(store.CloseAll) + client, server := yamuxPair() + store.Accept("peer-1", server) + + stream, err := client.OpenStream(GinkgoT().Context()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + var relayed net.Conn + Eventually(streams, "10s").Should(Receive(&relayed)) + go func() { + defer GinkgoRecover() + _, _ = stream.Write([]byte("hello")) + }() + buf := make([]byte, 5) + Expect(relayed.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = relayed.Read(buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("hello")) + }) + + It("replaces a peer's link when it dials again, and closes the one it lost", func() { + // A peer only re-dials because its previous link is gone from where it + // stands. Keeping both would leave a session nothing can be routed to, + // since the store holds one per peer. + store := cluster.NewSessionStore(nil) + DeferCleanup(store.CloseAll) + _, first := yamuxPair() + _, second := yamuxPair() + + store.Accept("peer-1", first) + store.Accept("peer-1", second) + + held, ok := store.Get("peer-1") + Expect(ok).To(BeTrue()) + Expect(held).To(BeIdenticalTo(second)) + Eventually(first.IsClosed, "10s").Should(BeTrue()) + Expect(second.IsClosed()).To(BeFalse(), "the link the peer is actually using was dropped") + }) + + It("forgets a session that ended, without evicting the one that replaced it", func() { + store := cluster.NewSessionStore(nil) + DeferCleanup(store.CloseAll) + client, server := yamuxPair() + store.Accept("peer-1", server) + + Expect(client.Close()).To(Succeed()) + Eventually(func() bool { + _, ok := store.Get("peer-1") + return ok + }, "10s").Should(BeFalse()) + }) + + It("closes every held link on shutdown, and refuses to store one afterwards", func() { + store := cluster.NewSessionStore(nil) + _, server := yamuxPair() + store.Accept("peer-1", server) + + store.CloseAll() + Eventually(server.IsClosed, "10s").Should(BeTrue()) + + _, late := yamuxPair() + store.Accept("peer-late", late) + _, ok := store.Get("peer-late") + Expect(ok).To(BeFalse()) + Eventually(late.IsClosed, "10s").Should(BeTrue(), + "a dial racing shutdown must not be left believing it holds a live link") + }) +}) diff --git a/core/services/cluster/wsconn_test.go b/core/services/cluster/wsconn_test.go index 305ce99b3c79..d7213fd87779 100644 --- a/core/services/cluster/wsconn_test.go +++ b/core/services/cluster/wsconn_test.go @@ -11,7 +11,6 @@ import ( "strings" "time" - clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" "github.com/mudler/LocalAI/core/services/cluster" "github.com/gorilla/websocket" @@ -239,7 +238,7 @@ var _ = Describe("Peer link payloads", func() { It("carries a payload far larger than one yamux frame end to end", func() { sessions := make(chan *yamux.Session, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) + servePeerRoute(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) srv := httptest.NewServer(e) DeferCleanup(srv.Close) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 72b957bcc27a..e147a995574d 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -443,14 +443,14 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - if err := db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}, &cluster.NodeConnection{}); err != nil { + if err := db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}); err != nil { return err } - // AutoMigrate models tables and columns but has no notion of a - // sequence, and the connection-ownership fence draws its epochs from - // one. It runs under this same lock so concurrently starting replicas - // do not race on the DDL. - return cluster.EnsureEpochSequence(context.Background(), db) + // The cluster package owns its own tables AND the sequence its + // ownership fence draws epochs from, which AutoMigrate cannot express. + // It runs under this same lock so concurrently starting replicas do not + // race on the DDL. + return cluster.Migrate(context.Background(), db) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 0231c2dc4a52..3a6f9768f669 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -64,6 +64,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 | | `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS server URL (e.g., `nats://localhost:4222`) | +| `--distributed-advertise-addr` | `LOCALAI_DISTRIBUTED_ADVERTISE_ADDR` | *(derived)* | `host:port` the **other frontend replicas** dial to reach this one. See [Replica peer links](#replica-peer-links). | | `--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) | | `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | **Umbrella switch.** Implies both `--nats-require-auth` and `--registration-require-auth` - one knob to lock down the NATS bus *and* the registration/file-transfer layer. Set this in production instead of the two granular flags. | @@ -78,6 +79,27 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | *(env only)* | `LOCALAI_MODEL_LOAD_WAIT` | `60s` | How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with `503`, a `Retry-After` header and live staging progress. The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to `0` to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front. See [Requests for a model that is still loading](#requests-for-a-model-that-is-still-loading). | | `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | +### Replica peer links + +Frontend replicas record themselves in an `instances` table and open direct links to each other, so that a request arriving at one replica can be served by state another replica holds. Each replica publishes one address for this, and every other replica dials it: it is the address **peers** use, which is not necessarily the address the process binds. A replica behind a Kubernetes Service, a load balancer or a NAT binds one and is reached at another. + +When `LOCALAI_DISTRIBUTED_ADVERTISE_ADDR` is unset, the address is derived: LocalAI asks the kernel which local address routes to PostgreSQL, and pairs it with the port it serves on. Every replica reaches the same database, so that address is on a network they demonstrably share. + +That only holds while the database is on **another host**. If PostgreSQL runs on the same host or pod (compose, single-node, a sidecar), the route to it is loopback, and advertising a loopback address would send every peer to itself. LocalAI refuses to guess in that case and logs: + +``` +This replica will not be reachable by its peers: no advertised address +``` + +The replica keeps serving every request that reaches it directly; what it cannot do is have another replica reach it. Set the address explicitly to fix it: + +```yaml +environment: + LOCALAI_DISTRIBUTED_ADVERTISE_ADDR: "10.0.1.7:8080" # or the pod IP, service DNS name, etc. +``` + +The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_REGISTRATION_TOKEN`, the same shared secret workers register with. Replicas that disagree about it cannot link. A replica that stops heartbeating for 30 seconds is dropped from the table by the others, along with the worker-connection rows it owned. + ### The model load deadline scales with the checkpoint The `LoadModel` deadline starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint read and pipeline init. That work is proportional to the bytes on disk, which makes any fixed deadline a model-size cliff rather than a timeout: a 70 GB video checkpoint on a Jetson Thor worker failed reproducibly against the old fixed 5m default (`rpc error: code = DeadlineExceeded` after 953.5s of wall clock, roughly 11m of which was backend install and staging), and simply raising the constant would only move the cliff to the next larger model while making a genuinely wedged *small* model hang for the whole inflated duration. diff --git a/tests/e2e/distributed/cluster/cluster.go b/tests/e2e/distributed/cluster/cluster.go index 1783d71e16f0..521882efe962 100644 --- a/tests/e2e/distributed/cluster/cluster.go +++ b/tests/e2e/distributed/cluster/cluster.go @@ -209,6 +209,14 @@ func (c *Cluster) startFrontend(i int, port int) (*Process, error) { // Pinning makes the cross-replica session a property of the harness. "LOCALAI_AUTH_HMAC_SECRET="+testHMACSecret, "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, + // Every replica here shares one host, so the address a peer dials is + // this process's own loopback address. It has to be said explicitly: + // the automatic discovery asks which local address routes to + // PostgreSQL, and this suite's PostgreSQL is a container published on + // 127.0.0.1, so the discovery refuses (correctly) rather than + // advertising a loopback address that would mean "yourself" on a + // multi-host deployment. + fmt.Sprintf("LOCALAI_DISTRIBUTED_ADVERTISE_ADDR=127.0.0.1:%d", port), "LOCALAI_AUTO_APPROVE_NODES=true", "DEBUG=true", ) @@ -334,6 +342,14 @@ func (c *Cluster) FrontendURL(i int) string { return fmt.Sprintf("http://127.0.0.1:%d", c.frontends[i].Port) } +// RegistrationToken is the shared secret this cluster was started with. It +// authenticates worker registration AND the replica-to-replica peer link, so a +// spec acting as a peer needs it rather than a second literal that can drift +// from Options. +func (c *Cluster) RegistrationToken() string { + return c.opts.RegistrationToken +} + // WorkerName is the node name worker i registered under. func (c *Cluster) WorkerName(i int) string { return c.workers[i].Name diff --git a/tests/e2e/distributed/cluster_baseline_test.go b/tests/e2e/distributed/cluster_baseline_test.go index 6cbaaa0ea915..63995b419156 100644 --- a/tests/e2e/distributed/cluster_baseline_test.go +++ b/tests/e2e/distributed/cluster_baseline_test.go @@ -122,6 +122,14 @@ func mockBackendBinary() string { // existing caller keeps the plain two-argument form and the default shape. func startCluster(frontends, workers int, customise ...func(*cluster.Options)) *cluster.Cluster { GinkgoHelper() + c, _ := startClusterOnFreshDB(frontends, workers, customise...) + return c +} + +// startClusterOnFreshDB is startCluster plus the DSN of the database it was +// given, for a spec that has to read a table no endpoint exposes. +func startClusterOnFreshDB(frontends, workers int, customise ...func(*cluster.Options)) (*cluster.Cluster, string) { + GinkgoHelper() // Resolved before SetupInfra so a missing binary skips without having paid // for a database that the skip would then leave to DeferCleanup. @@ -163,7 +171,7 @@ func startCluster(frontends, workers int, customise ...func(*cluster.Options)) * } c.Stop() }) - return c + return c, infra.PGURL } // rosterProbe polls one frontend's node roster. diff --git a/tests/e2e/distributed/cluster_peerlink_test.go b/tests/e2e/distributed/cluster_peerlink_test.go new file mode 100644 index 000000000000..8055ba297a67 --- /dev/null +++ b/tests/e2e/distributed/cluster_peerlink_test.go @@ -0,0 +1,263 @@ +package distributed_test + +import ( + "context" + "fmt" + "net" + "strings" + "time" + + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/postgres" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" +) + +const ( + // instanceRosterTimeout bounds the wait for a replica's row to appear. + // Registration is synchronous in startup, so this only has to cover the gap + // between /readyz answering and this spec's first query. + instanceRosterTimeout = "30s" + instanceRosterPoll = "500ms" + + // deadReplicaTimeout bounds the wait for a survivor to reap a replica that + // was killed: the liveness window plus a sweep interval plus slack. It is + // deliberately derived from the constants rather than a round number, so + // tightening the window shortens the spec instead of leaving it passing for + // the wrong reason. + deadReplicaTimeout = clustersvc.InstanceLiveness + 4*clustersvc.InstanceHeartbeat + + // peerDialTimeout bounds one peer dial. Every replica here is a local + // process, so a dial that needs longer has failed, not slowed. + peerDialTimeout = 20 * time.Second +) + +// openClusterDB connects to the database the cluster was given, so a spec can +// read the tables the peer link keeps. Nothing serves them over HTTP: they are +// replica-to-replica state, not an admin surface, and inventing an endpoint to +// observe them would be a bigger change than the thing under test. +func openClusterDB(dsn string) *gorm.DB { + GinkgoHelper() + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: gormlogger.Discard}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { closeDB(db) }) + return db +} + +// hostPortOf strips the scheme off a frontend URL, giving the form the +// instances table stores. +func hostPortOf(url string) string { + return strings.TrimPrefix(strings.TrimPrefix(url, "http://"), "https://") +} + +// instanceRoster reads the live replica rows, keeping the last error so a +// failing Eventually can name it. +type instanceRoster struct { + registry *clustersvc.Registry + ctx context.Context + + lastErr error + lastSaw []clustersvc.Instance +} + +func newInstanceRoster(db *gorm.DB) *instanceRoster { + return &instanceRoster{registry: clustersvc.NewRegistry(db), ctx: context.Background()} +} + +// addresses returns the advertised address of every live replica, or nil on a +// query error so Eventually keeps trying. +func (r *instanceRoster) addresses() []string { + live, err := r.registry.Live(r.ctx, clustersvc.InstanceLiveness) + if err != nil { + r.lastErr = err + return nil + } + r.lastErr = nil + r.lastSaw = live + addrs := []string{} + for _, instance := range live { + addrs = append(addrs, instance.AdvertisedAddr) + } + return addrs +} + +// idAt returns the id of the live replica advertising addr, or "" if no such +// row is present yet. +func (r *instanceRoster) idAt(addr string) string { + for _, instance := range r.lastSaw { + if instance.AdvertisedAddr == addr { + return instance.ID + } + } + return "" +} + +func (r *instanceRoster) describe() string { + if r.lastErr != nil { + return fmt.Sprintf("the last read of the instances table failed: %v", r.lastErr) + } + return fmt.Sprintf("the instances table held %d live replica(s): %+v", len(r.lastSaw), r.lastSaw) +} + +// awaitReplicas waits for every frontend of c to publish its address and +// returns the roster, positioned on that reading. +func awaitReplicas(roster *instanceRoster, addrs ...string) { + GinkgoHelper() + Eventually(roster.addresses, instanceRosterTimeout, instanceRosterPoll). + Should(ConsistOf(addrs), roster.describe) +} + +var _ = Describe("Cluster peer link", Label("Distributed"), Label("Cluster"), func() { + It("publishes an address for every replica that peers can actually dial", func() { + // A wrong implementation registers nothing (the whole of phase 1 had no + // call site until this spec), registers one row for two replicas, or + // records an address nothing can connect to: the bind address of a + // replica behind a service, or the loopback address the route to a + // co-located database would suggest. + c, dsn := startClusterOnFreshDB(2, 0) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + + // "Routable" is not a property of the string. Connect to each address, + // which is the only check that would have caught a replica publishing + // the port it was configured with rather than the one it serves on. + for _, instance := range roster.lastSaw { + conn, err := net.DialTimeout("tcp", instance.AdvertisedAddr, peerDialTimeout) + Expect(err).ToNot(HaveOccurred(), + "replica %s advertises %q, which nothing can connect to", instance.ID, instance.AdvertisedAddr) + Expect(conn.Close()).To(Succeed()) + } + }) + + It("carries a peer stream between two replicas, and refuses one without the cluster token", func() { + // A wrong implementation fails here on WebSocket framing, which is the + // likeliest defect in the peer link: the adapter has to turn + // message-oriented WebSocket frames into the undelimited byte stream + // yamux drives. It also fails if the route was never registered on the + // real server, or if the global session middleware answers it: a peer + // carries no session and no user, only the cluster token. + // + // The stream is opened with the production dialler, resolving the peer + // through the production registry, over a real socket to a real + // process. This spec plays the sibling replica, because phase 1 has + // nothing that makes a frontend dial one on its own. + c, dsn := startClusterOnFreshDB(2, 0) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + + peerID := roster.idAt(hostPortOf(c.FrontendURL(1))) + Expect(peerID).ToNot(BeEmpty()) + + ctx, cancel := context.WithTimeout(context.Background(), peerDialTimeout) + defer cancel() + + pool := clustersvc.NewPeerPool("e2e-peer", c.RegistrationToken(), roster.registry) + DeferCleanup(pool.Close) + + // OpenStream is only acknowledged once the far side accepts, so this + // returning at all proves the frontend is accepting streams on the + // session it took, in addition to proving the handshake. + stream, err := pool.Open(ctx, peerID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + // Phase 1 installs no relay, so the accepted stream is refused at once. + // What matters here is that the refusal arrives: a replica that held + // the session without accepting on it would leave this read parked + // until the deadline, which is the failure mode a live cluster would + // experience as every relayed request hanging. + Expect(stream.SetReadDeadline(time.Now().Add(peerDialTimeout))).To(Succeed()) + _, err = stream.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred(), + "the peer accepted the stream and then neither answered nor closed it") + Expect(err).ToNot(MatchError(context.DeadlineExceeded)) + + // The same dial with the wrong credentials must be refused, otherwise + // the success above says nothing about authentication. + impostor := clustersvc.NewPeerPool("e2e-peer", "not-the-cluster-token", roster.registry) + DeferCleanup(impostor.Close) + _, err = impostor.Open(ctx, peerID) + Expect(err).To(MatchError(clustersvc.ErrPeerUnreachable)) + Expect(err).ToNot(MatchError(clustersvc.ErrInstanceNotFound), + "a peer refusing credentials is a live peer; reading it as absence is how a replica evicts healthy workers") + }) + + It("reports a killed replica as unreachable, reaps what it owned, and evicts no worker", func() { + // This is the absence rule, pinned before phase 2 can depend on it. A + // wrong implementation lets a peer that will not answer surface as node + // absence, and a caller entitled to act on absence then reclaims what + // the peer was running: a network hiccup between two healthy replicas + // evicts healthy workers. + // + // It also pins the reaper: the connection rows a dead replica owned are + // swept by the same sweeper that decides the replica is dead, so the + // two can never disagree about who is alive. + c, dsn := startClusterOnFreshDB(2, 1) + + client, err := c.AdminSession(0) + Expect(err).ToNot(HaveOccurred()) + + // The worker registers with frontend 0, so frontend 1 is the replica + // that can die without taking the worker's registrar with it. + registrar, err := c.WorkerRegistrar(0) + Expect(err).ToNot(HaveOccurred()) + Expect(registrar).To(Equal(0), "this spec kills frontend 1 and needs the worker to have registered elsewhere") + + probe := newRosterProbe(c, client, 0) + Eventually(probe.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ContainElement(c.WorkerName(0)), probe.describe) + workerID := probe.idOf(c.WorkerName(0)) + Expect(workerID).ToNot(BeEmpty()) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + survivorID := roster.idAt(hostPortOf(c.FrontendURL(0))) + doomedID := roster.idAt(hostPortOf(c.FrontendURL(1))) + Expect(survivorID).ToNot(BeEmpty()) + Expect(doomedID).ToNot(BeEmpty()) + + // Give frontend 1 the worker's tunnel. Phase 2 makes the worker do this + // by dialling; here the claim is written directly, because the point + // under test is what happens to the claim when its owner dies. + ctx := context.Background() + epoch, err := roster.registry.Claim(ctx, workerID, doomedID) + Expect(err).ToNot(HaveOccurred()) + Expect(epoch).ToNot(BeZero()) + + Expect(c.KillFrontend(1)).To(Succeed()) + Eventually(func() bool { return c.FrontendAlive(1) }, "20s", "500ms").Should(BeFalse()) + + // The row is still there for the whole liveness window, so this is the + // case that matters: the peer is KNOWN and will not answer. + dialCtx, cancel := context.WithTimeout(ctx, peerDialTimeout) + defer cancel() + pool := clustersvc.NewPeerPool("e2e-peer", c.RegistrationToken(), roster.registry) + DeferCleanup(pool.Close) + _, err = pool.Open(dialCtx, doomedID) + Expect(err).To(MatchError(clustersvc.ErrPeerUnreachable)) + Expect(err).ToNot(MatchError(clustersvc.ErrInstanceNotFound), + "a dead replica whose row is still present is unreachable, not absent") + + // The survivor sweeps the dead replica and, in the same pass, the claim + // it left behind. + Eventually(roster.addresses, deadReplicaTimeout, instanceRosterPoll). + Should(ConsistOf(hostPortOf(c.FrontendURL(0))), roster.describe) + ownerErr := func() error { + _, _, err := roster.registry.Owner(ctx, workerID) + return err + } + Eventually(ownerErr, deadReplicaTimeout, instanceRosterPoll). + Should(MatchError(clustersvc.ErrNoConnection), + "the claim held by a replica that no longer exists was never reaped") + + // And the worker itself is untouched throughout. Nothing about a peer + // dying may reach the node roster. + Consistently(probe.healthyNames, "6s", "1s"). + Should(ContainElement(c.WorkerName(0)), probe.describe) + }) +}) From e26d556594a2596e115d59daad14ab97e3fe59f2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 03:08:33 +0000 Subject: [PATCH 15/42] fix(cluster): hold the guarantees phase 1's comments were claiming Review found the recurring class: assertions that a wrong implementation also satisfies. The "refuse promptly, never park the peer" guarantee was stated in three places and tested in none. Removing the Close from the no-relay branch left the whole cluster suite green, because the specs asserted only that some error arrived and yamux reports a read deadline as ErrTimeout: a parked stream satisfied that as well as a refused one. Both specs now require an ENDING, EOF or a reset, inside a deadline short enough that parking is unmistakable, and both go red when the Close is removed. Deregistration existed only in a comment. Membership.Stop ended the loop and left the row behind, so every clean rolling restart had peers dialling a corpse for the full liveness window; the shutdown comment described the opposite. Registry.Deregister deletes the row and the connections that replica owned, in one transaction, for the reason the sweeper does both, and an e2e spec pins departure inside a budget shorter than the liveness window so it cannot pass on the sweeper doing the work. Before: the spec times out with both replicas still live. After: 3.6s. The configured advertised address bypassed every check discovery makes, so the one value most likely to be copied between hosts, 127.0.0.1, was taken verbatim and would make every peer dial itself. Both paths now share one rejection rule: unparseable is refused, "this host" is warned about once and honoured, because a single-host deployment uses it correctly. Two comments claimed more than the code does. The sweeper said a stalled replica recovers via re-register; only its instance row does, while the connections another replica reaped stay gone and the sockets stay held here - phase 2 must re-claim, on re-register, every connection a replica still holds locally. And Owner became OwnerRow, documenting that the owner it names may be dead for up to InstanceLiveness plus a heartbeat and that any caller acting on it must join instances itself, so the deferred constraint lives at the call site rather than in a report; the plain name is left free for the joining version. Minors: warn once when the peer link mounts with no registration token, so an operator sees the cause rather than 401s; Stop no longer blocks forever when Start was never called; corrected the NewRegistry migration doc and an e2e comment that described a 6s window as "throughout". Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 17 +++- core/http/app.go | 9 ++ core/services/cluster/instance.go | 66 +++++++++++++-- core/services/cluster/instance_test.go | 46 ++++++++++ core/services/cluster/membership.go | 84 ++++++++++++++++++- core/services/cluster/membership_test.go | 69 ++++++++++++++- core/services/cluster/ownership.go | 19 ++++- core/services/cluster/ownership_test.go | 12 +-- core/services/cluster/sessions_test.go | 16 +++- .../e2e/distributed/cluster_peerlink_test.go | 72 +++++++++++++--- 10 files changed, 369 insertions(+), 41 deletions(-) diff --git a/core/application/distributed.go b/core/application/distributed.go index 404b2112c3fc..8769f006d228 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -509,8 +509,21 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // is only a peer-reachable answer when the database is on another host; // DiscoverAdvertisedAddr refuses rather than guessing when it is not. func advertisedPeerAddr(cfg *config.ApplicationConfig) (string, error) { - if cfg.Distributed.AdvertiseAddr != "" { - return cfg.Distributed.AdvertiseAddr, nil + if configured := cfg.Distributed.AdvertiseAddr; configured != "" { + // A configured address skips discovery, so it also skips every check + // discovery makes. Unusable is refused; merely questionable (a + // loopback address, correct on one host and wrong on three) is said + // once and honoured, because refusing it would refuse single-host + // deployments that use it correctly. + reason, err := cluster.CheckAdvertisedAddr(configured) + if err != nil { + return "", err + } + if reason != "" { + xlog.Warn("Configured peer address is not one another host can dial", + "address", configured, "reason", reason, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + } + return configured, nil } if cfg.APIAddress == "" { return "", fmt.Errorf("no API address to derive a peer port from") diff --git a/core/http/app.go b/core/http/app.go index d6418165b5f0..a2a3555eafe0 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" + clustersvc "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" @@ -581,6 +582,14 @@ func API(application *application.Application) (*echo.Echo, error) { // registration token, so publishing it unconditionally would put a // multiplexer on every single-binary install. if d := application.Distributed(); d != nil && d.PeerSessions != nil { + if distCfg.RegistrationToken == "" { + // The handler fails closed on an empty token, which is right and + // invisible: without this line an operator sees only 401s on a + // route they never configured, and nothing connecting them to the + // token they did not set. + xlog.Warn("Replica peer link will refuse every dial: no registration token is configured", + "route", clustersvc.PeerPath, "knob", "LOCALAI_REGISTRATION_TOKEN") + } routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, d.PeerSessions.Accept) } diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index 027b0db3b0d9..d7e7afa71740 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -38,9 +38,10 @@ type Registry struct { db *gorm.DB } -// NewRegistry returns a Registry over db. Migration is the caller's job; the -// nodes registry owns the AutoMigrate for every table in this deployment so -// that a single advisory lock covers them all. +// NewRegistry returns a Registry over db. Migration is the caller's job: this +// package's tables and sequence are created by Migrate, which the nodes +// registry calls under the one advisory lock that covers every table in the +// deployment. func NewRegistry(db *gorm.DB) *Registry { return &Registry{db: db} } @@ -152,20 +153,67 @@ func DiscoverAdvertisedAddr(dsn string, port int) (string, error) { // information about the address we just read. defer func() { _ = conn.Close() }() local, ok := conn.LocalAddr().(*net.UDPAddr) - if !ok || local.IP == nil || local.IP.IsUnspecified() { + 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 local.IP.IsLoopback() { - return "", fmt.Errorf("the route to database host %q is loopback (%s), so the database is local to this replica and its peer-reachable address cannot be discovered; set the advertised address explicitly", host, local.IP) + 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) // A zone is only ever attached to a scoped (link-local) address, so this is // the same rejection stated twice; the Zone check keeps the guarantee if a // platform ever hands back a scoped address of another class, because // IP.String() would silently drop the %iface and yield an undialable host. - if local.IP.IsLinkLocalUnicast() || local.Zone != "" { - return "", fmt.Errorf("the route to database host %q is link-local (%s), which peers on other hosts cannot dial; set the advertised address explicitly", host, local.IP) + case ip.IsLinkLocalUnicast() || zone != "": + return fmt.Sprintf("link-local (%s), which peers on other hosts cannot dial", ip) } - return net.JoinHostPort(local.IP.String(), strconv.Itoa(port)), nil + return "" +} + +// 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) + } + // 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, ""), nil } // dsnHostPort extracts the host and port from either DSN form gorm's postgres diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go index d8075a7a3cad..d27043e03c46 100644 --- a/core/services/cluster/instance_test.go +++ b/core/services/cluster/instance_test.go @@ -113,3 +113,49 @@ var _ = Describe("Advertised address discovery", func() { 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")) + }) +}) diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index 3b1182b17458..f3f71a321acd 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -25,6 +25,10 @@ const ( // owner is merely slow has to be re-homed for nothing. The cost of waiting // is bounded and symmetric: traffic for that worker is retried, not lost. InstanceLiveness = 30 * time.Second + + // deregisterTimeout bounds the deregistration Stop performs. Shutdown is + // not the place to wait on a database. + deregisterTimeout = 5 * time.Second ) // Membership publishes this replica's address and keeps the instances table @@ -45,6 +49,10 @@ type Membership struct { stop chan struct{} done chan struct{} stopOnce sync.Once + + // mu guards started, which tells Stop whether there is a loop to join. + mu sync.Mutex + started bool } // NewMembership returns the membership loop for one replica. The address is @@ -72,17 +80,49 @@ func (m *Membership) Start(ctx context.Context) error { return err } xlog.Info("Cluster instance registered", "id", m.id, "addr", m.addr) + m.mu.Lock() + m.started = true + m.mu.Unlock() go m.loop(ctx) return nil } -// Stop ends the loop and waits for it. Safe to call more than once. +// Stop ends the loop, waits for it, and removes this replica's row. +// +// Deregistering is what makes a rolling restart quick for everyone else: a +// replica that just closes its sockets is indistinguishable from one that +// crashed, so its peers keep dialling it for the whole liveness window. It is +// best-effort by nature (a killed process never gets here), which is why the +// sweeper still exists. +// +// Safe to call more than once, and on a Membership that was never started. func (m *Membership) Stop() { if m == nil { return } - m.stopOnce.Do(func() { close(m.stop) }) - <-m.done + m.mu.Lock() + started := m.started + m.mu.Unlock() + if started { + m.stopOnce.Do(func() { close(m.stop) }) + // Only a started Membership ever closes done. Waiting on one that was + // never started, or whose Start failed, would block forever. + <-m.done + } + + // Deliberately NOT the context Start was given: that one is the + // application's, and by the time anything calls Stop it has usually been + // cancelled already, so deregistering on it would fail every time. The + // bound is here instead, because shutdown must not hang on a database that + // went away before the process using it. + 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 will 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) { @@ -114,6 +154,10 @@ func (m *Membership) tick(ctx context.Context) { // Another replica swept this row while this process was stalled long // enough to look dead. Re-register rather than heartbeat: a heartbeat // carries no address, so the row has to be rebuilt from scratch. + // + // This rebuilds the instance row ONLY. The sweep that removed it also + // removed the connections this replica owned, and re-claiming those + // needs the tunnel registry phase 2 introduces; see ReapStale. xlog.Warn("Cluster instance row was reaped, re-registering", "id", m.id) if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { xlog.Error("Re-registering cluster instance failed", "id", m.id, "error", err) @@ -132,6 +176,30 @@ func (m *Membership) tick(ctx context.Context) { } } +// Deregister removes one replica and the connections it owned. +// +// It deletes both, in one transaction, for the same reason ReapStale does: a +// replica that is gone owns nothing, and leaving its connection rows behind +// would point every reader at an owner that no longer exists. This is the +// announced form of what the sweeper does by inference, and the two must not +// disagree about what "gone" removes. +func (r *Registry) Deregister(ctx context.Context, id string) error { + if err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := tx.Where("owner_instance_id = ?", id).Delete(&NodeConnection{}).Error; err != nil { + return fmt.Errorf("deleting connections owned by %q: %w", id, err) + } + // No RowsAffected check: deregistering a row another replica already + // swept is the normal outcome of a slow shutdown, not an error. + if err := tx.Where("id = ?", id).Delete(&Instance{}).Error; err != nil { + return fmt.Errorf("deleting instance %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 the connection rows whose owner is no longer among the survivors. // @@ -144,7 +212,15 @@ func (m *Membership) tick(ctx context.Context) { // self is never reaped. This process may fail to heartbeat for longer than the // window (a long stall, a database blip) and still be serving: deleting its own // row would then delete the connections of workers that are, at that moment, -// connected to it. The stall is recovered by the re-register in tick instead. +// connected to it. +// +// That protection is one-sided, and only the instance row recovers on its own. +// A replica that stalls long enough is reaped BY ANOTHER replica, taking its +// connection rows with it, and the re-register in tick rebuilds the instance +// row and nothing else: the sockets are still held here while the table says +// nobody holds them. Phase 2 closes this by re-claiming, on re-register, every +// connection this replica still holds locally, which needs the tunnel registry +// that owns those sockets. // // PostgreSQL only, like Live: distributed mode requires it, and the interval // arithmetic is measured on the database's clock because liveness is compared diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go index a00a08c0b954..d490d529dd47 100644 --- a/core/services/cluster/membership_test.go +++ b/core/services/cluster/membership_test.go @@ -47,7 +47,7 @@ var _ = Describe("Reaping dead replicas", func() { Expect(connections).To(Equal(int64(1)), "a worker whose owner no longer exists is recorded as connected to nothing") - _, _, err = reg.Owner(ctx, "w1") + _, _, err = reg.OwnerRow(ctx, "w1") Expect(err).To(MatchError(cluster.ErrNoConnection)) }) @@ -61,7 +61,7 @@ var _ = Describe("Reaping dead replicas", func() { Expect(err).ToNot(HaveOccurred()) Expect(connections).To(BeZero()) - owner, stored, err := reg.Owner(ctx, "w1") + owner, stored, err := reg.OwnerRow(ctx, "w1") Expect(err).ToNot(HaveOccurred()) Expect(owner).To(Equal("other")) Expect(stored).To(Equal(epoch)) @@ -82,11 +82,74 @@ var _ = Describe("Reaping dead replicas", func() { Expect(instances).To(BeZero()) Expect(connections).To(BeZero()) - owner, _, err := reg.Owner(ctx, "w1") + owner, _, err := reg.OwnerRow(ctx, "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", "10.0.0.2:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "staying", "10.0.0.1:8080", "v1")).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 = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + owner, _, err := reg.OwnerRow(ctx, "w2") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("staying"), "deregistering one replica took another replica's claim") + }) + + It("deregisters when the membership loop stops", func() { + membership := cluster.NewMembership(reg, "me", "10.0.0.1:8080", "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("tolerates a repeated deregistration, because a sweeper may have got there first", func() { + Expect(reg.Register(ctx, "gone", "10.0.0.2:8080", "v1")).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() { + // Nothing calls this today. It exists because the loop channel is only + // ever closed by a started loop, so joining an unstarted one blocks + // forever, and phase 2 adds callers to this shutdown path. + membership := cluster.NewMembership(reg, "never-started", "10.0.0.1:8080", "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", "10.0.0.2:8080", "v1")).To(Succeed()) age("dead", time.Hour) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index 3b6b6cb622e4..131b3d0b976d 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -161,9 +161,22 @@ func (r *Registry) Claim(ctx context.Context, nodeID, ownerID string) (int64, er return epoch, nil } -// Owner returns the replica that holds nodeID's tunnel and the epoch of that -// claim, or ErrNoConnection when the node has no recorded connection. -func (r *Registry) Owner(ctx context.Context, nodeID string) (string, int64, error) { +// OwnerRow returns the row recording which replica holds nodeID's tunnel, and +// the epoch of that claim, or ErrNoConnection when the node has no recorded +// connection. +// +// It answers "what does the table say", NOT "who holds this tunnel". The owner +// it names may be dead: a replica that dies stops heartbeating, and its rows +// survive until another replica's sweep removes them, which is up to +// InstanceLiveness plus one InstanceHeartbeat later. Any caller that ACTS on +// the answer must join instances itself and treat a non-live owner as +// ErrNoConnection; relaying to the row without that check is relaying into a +// process that is gone. +// +// The name says row on purpose, so that the joining version can take the plain +// name when phase 2 introduces the first caller that needs it. Nothing in +// phase 1 reads this outside tests, which is why the join is not here yet. +func (r *Registry) OwnerRow(ctx context.Context, nodeID string) (string, int64, error) { var conn NodeConnection err := r.db.WithContext(ctx).Where("node_id = ?", nodeID).First(&conn).Error if errors.Is(err, gorm.ErrRecordNotFound) { diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index 5744e07cd170..6d6fdf58bc7e 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -85,14 +85,14 @@ var _ = Describe("Connection ownership", func() { e2, err := reg.Claim(ctx, "w1", "inst-b") Expect(err).ToNot(HaveOccurred()) - owner, epoch, err := reg.Owner(ctx, "w1") + owner, epoch, err := reg.OwnerRow(ctx, "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 := reg.Owner(ctx, "ghost") + _, _, err := reg.OwnerRow(ctx, "ghost") Expect(err).To(MatchError(cluster.ErrNoConnection)) }) @@ -105,7 +105,7 @@ var _ = Describe("Connection ownership", func() { // inst-a tries to clean up after losing the claim. Expect(reg.Release(ctx, "w1", "inst-a", e1)).ToNot(Succeed()) - owner, _, err := reg.Owner(ctx, "w1") + owner, _, err := reg.OwnerRow(ctx, "w1") Expect(err).ToNot(HaveOccurred()) Expect(owner).To(Equal("inst-b"), "a stale owner must not be able to delete a live claim") }) @@ -120,7 +120,7 @@ var _ = Describe("Connection ownership", func() { Expect(reg.Release(ctx, "w1", "inst-a", e1)).ToNot(Succeed()) - owner, _, err := reg.Owner(ctx, "w1") + owner, _, err := reg.OwnerRow(ctx, "w1") Expect(err).ToNot(HaveOccurred()) Expect(owner).To(Equal("inst-a")) }) @@ -144,7 +144,7 @@ var _ = Describe("Connection ownership", func() { // live claim disappearing rather than on the epoch arithmetic. Expect(reg.Release(ctx, "w1", "inst-a", eA1)).ToNot(Succeed()) - owner, epoch, err := reg.Owner(ctx, "w1") + owner, epoch, err := reg.OwnerRow(ctx, "w1") Expect(err).ToNot(HaveOccurred(), "the delayed cleanup deleted the live claim") Expect(owner).To(Equal("inst-a")) Expect(epoch).To(Equal(eA2)) @@ -156,7 +156,7 @@ var _ = Describe("Connection ownership", func() { Expect(err).ToNot(HaveOccurred()) Expect(reg.Release(ctx, "w1", "inst-a", e)).To(Succeed()) - _, _, err = reg.Owner(ctx, "w1") + _, _, err = reg.OwnerRow(ctx, "w1") Expect(err).To(MatchError(cluster.ErrNoConnection)) }) diff --git a/core/services/cluster/sessions_test.go b/core/services/cluster/sessions_test.go index f3fb44ca8cae..29ca9b537d42 100644 --- a/core/services/cluster/sessions_test.go +++ b/core/services/cluster/sessions_test.go @@ -1,6 +1,7 @@ package cluster_test import ( + "io" "net" "time" @@ -29,6 +30,11 @@ func yamuxPair() (client *yamux.Session, server *yamux.Session) { return client, server } +// refusalDeadline bounds how long a refused stream may take to end. A refusal +// is one frame from a peer that already decided, so anything near this is the +// hang it exists to detect. +const refusalDeadline = 2 * time.Second + var _ = Describe("Accepted peer sessions", func() { It("accepts and refuses a stream rather than leaving the peer parked", func() { // yamux only acknowledges a stream once the far side accepts it, so a @@ -43,9 +49,15 @@ var _ = Describe("Accepted peer sessions", func() { Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = stream.Close() }) - Expect(stream.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + // The deadline is short and is NOT the thing being asserted: yamux + // reports a deadline as ErrTimeout, and requiring an ending instead + // (EOF from the peer's Close, or a reset) is what separates "refused" + // from "parked". An earlier version asserted only that some error + // arrived, which a parked stream satisfies just as well. + Expect(stream.SetReadDeadline(time.Now().Add(refusalDeadline))).To(Succeed()) _, err = stream.Read(make([]byte, 1)) - Expect(err).To(HaveOccurred(), "a refused stream must end, not hang") + Expect(err).To(SatisfyAny(MatchError(io.EOF), MatchError(yamux.ErrStreamReset)), + "a refused stream must END within %s; %v means the peer accepted it and then left it parked", refusalDeadline, err) }) It("hands a stream to the relay when one is installed", func() { diff --git a/tests/e2e/distributed/cluster_peerlink_test.go b/tests/e2e/distributed/cluster_peerlink_test.go index 8055ba297a67..d99fcf70cfd8 100644 --- a/tests/e2e/distributed/cluster_peerlink_test.go +++ b/tests/e2e/distributed/cluster_peerlink_test.go @@ -3,12 +3,15 @@ package distributed_test import ( "context" "fmt" + "io" "net" "strings" "time" clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "gorm.io/driver/postgres" @@ -33,6 +36,17 @@ const ( // peerDialTimeout bounds one peer dial. Every replica here is a local // process, so a dial that needs longer has failed, not slowed. peerDialTimeout = 20 * time.Second + + // gracefulDepartureTimeout bounds the wait for a cleanly stopped replica to + // leave the table. It must stay well under InstanceLiveness, which the spec + // asserts: a budget that reached the window would pass on the sweeper doing + // the work and prove nothing about deregistration. + gracefulDepartureTimeout = 15 * time.Second + + // peerRefusalTimeout bounds how long a refused stream may take to end. It + // is short on purpose: the refusal is one frame from a replica that has + // already decided, so a stream still open at this point is parked. + peerRefusalTimeout = 5 * time.Second ) // openClusterDB connects to the database the cluster was given, so a spec can @@ -166,16 +180,16 @@ var _ = Describe("Cluster peer link", Label("Distributed"), Label("Cluster"), fu Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = stream.Close() }) - // Phase 1 installs no relay, so the accepted stream is refused at once. - // What matters here is that the refusal arrives: a replica that held - // the session without accepting on it would leave this read parked - // until the deadline, which is the failure mode a live cluster would - // experience as every relayed request hanging. - Expect(stream.SetReadDeadline(time.Now().Add(peerDialTimeout))).To(Succeed()) + // Phase 1 installs no relay, so the accepted stream must be refused at + // once: an ENDING (EOF from the peer's Close, or a reset), not merely + // an error. A replica that accepted the stream and then left it parked + // would fail this read too, but with yamux's ErrTimeout, and that is + // the failure a live cluster experiences as every relayed request + // hanging until its own deadline. + Expect(stream.SetReadDeadline(time.Now().Add(peerRefusalTimeout))).To(Succeed()) _, err = stream.Read(make([]byte, 1)) - Expect(err).To(HaveOccurred(), - "the peer accepted the stream and then neither answered nor closed it") - Expect(err).ToNot(MatchError(context.DeadlineExceeded)) + Expect(err).To(SatisfyAny(MatchError(io.EOF), MatchError(yamux.ErrStreamReset)), + "the peer accepted the stream and then neither answered nor ended it: %v", err) // The same dial with the wrong credentials must be refused, otherwise // the success above says nothing about authentication. @@ -187,6 +201,38 @@ var _ = Describe("Cluster peer link", Label("Distributed"), Label("Cluster"), fu "a peer refusing credentials is a live peer; reading it as absence is how a replica evicts healthy workers") }) + It("stops being dialled as soon as a replica shuts down cleanly", func() { + // The crash case below is handled by the sweeper, at the cost of a + // whole liveness window of peers dialling a corpse. A rolling update is + // not a crash: the replica knows it is leaving and says so. Without + // deregistration the two are indistinguishable, and every rolling + // restart spends that window failing peer dials for no reason. + c, dsn := startClusterOnFreshDB(2, 0) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + departingID := roster.idAt(hostPortOf(c.FrontendURL(1))) + Expect(departingID).ToNot(BeEmpty()) + + Expect(c.StopFrontendGracefully(1)).To(Succeed()) + Eventually(func() bool { return c.FrontendAlive(1) }, "20s", "500ms").Should(BeFalse()) + + // The budget is deliberately shorter than the liveness window: passing + // it proves the replica announced its departure rather than aged out. + Expect(gracefulDepartureTimeout).To(BeNumerically("<", clustersvc.InstanceLiveness)) + Eventually(roster.addresses, gracefulDepartureTimeout, instanceRosterPoll). + Should(ConsistOf(hostPortOf(c.FrontendURL(0))), roster.describe) + + // And absence is the RIGHT answer here, unlike the killed case: the + // replica said it was going. A caller may act on this. + ctx, cancel := context.WithTimeout(context.Background(), peerDialTimeout) + defer cancel() + pool := clustersvc.NewPeerPool("e2e-peer", c.RegistrationToken(), roster.registry) + DeferCleanup(pool.Close) + _, err := pool.Open(ctx, departingID) + Expect(err).To(MatchError(clustersvc.ErrInstanceNotFound)) + }) + It("reports a killed replica as unreachable, reaps what it owned, and evicts no worker", func() { // This is the absence rule, pinned before phase 2 can depend on it. A // wrong implementation lets a peer that will not answer surface as node @@ -248,15 +294,17 @@ var _ = Describe("Cluster peer link", Label("Distributed"), Label("Cluster"), fu Eventually(roster.addresses, deadReplicaTimeout, instanceRosterPoll). Should(ConsistOf(hostPortOf(c.FrontendURL(0))), roster.describe) ownerErr := func() error { - _, _, err := roster.registry.Owner(ctx, workerID) + _, _, err := roster.registry.OwnerRow(ctx, workerID) return err } Eventually(ownerErr, deadReplicaTimeout, instanceRosterPoll). Should(MatchError(clustersvc.ErrNoConnection), "the claim held by a replica that no longer exists was never reaped") - // And the worker itself is untouched throughout. Nothing about a peer - // dying may reach the node roster. + // And the worker survives the sweep that removed its owner. This is a + // window after the reaping, not a watch over the whole scenario: + // Consistently starts here, so what it rules out is the sweep, or + // anything reacting to it, taking the worker with it. Consistently(probe.healthyNames, "6s", "1s"). Should(ContainElement(c.WorkerName(0)), probe.describe) }) From 0d13056d53cf24965cddef1e12b5eed0308b8cce Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 03:26:36 +0000 Subject: [PATCH 16/42] fix(cluster): share one lock order, and correct the phase 1 comments ReapStale deleted from instances then node_connections while Deregister took them the other way round, both inside one transaction and both running concurrently by design: a replica shuts down while a peer sweeps it. Opposite orders let each hold the row the other waits for. PostgreSQL breaks the cycle by aborting one side, so the cost today is a warning rather than lost data, but the inversion costs nothing to remove. Deregister now deletes the instance row first. That is the order ReapStale is forced into anyway, since its connection delete asks which instance rows survived, so the sweeper is the fixed side. Both functions say the order is deliberate and shared, and name the other. A spec records the statements each path issues and asserts they delete from the same two tables in the same order; racing two transactions until they really deadlock would be flaky and could pass for the wrong reason. The rest is comment and spec accuracy, deferred from the phase 1 task reviews: - co-location does not imply loopback. Compose's usual host=postgres resolves to a bridge address and discovery works there; it is a DSN that NAMES localhost that yields a loopback source address. Corrected in the DiscoverAdvertisedAddr doc and in the spec comment that repeated it. - unroutableReason labelled every scoped address "link-local", including the class the check exists for, and formatted the IP with %s, which drops the %iface, so the reported address was not the one being rejected. Split into two cases, both rendered with their zone. CheckAdvertisedAddr passed zone "" and net.ParseIP rejects fe80::1%eth0, so a scoped literal looked like a name and collected no warning at all; the zone is now split off before parsing. - Splice's "Both callers satisfy it" claimed callers that still do not exist. It now names the two stream types the wake-on-Close property was verified against and says a phase 2 caller over anything else has to check it. - restored, short, why a socket-level ECONNRESET stays reported while a yamux reset does not: the yamux endings are the teardown Splice's own Close provokes, and whether an aborted request is routine is the relay's policy. - the real-yamux spec's far.Read had no deadline, so a stall parked the suite rather than failing it. - gorilla's SetWriteDeadline is conn.go:796, not 787. - ClusterPathPrefix is no longer derived from: the peer route spells its path out, because core/services/cluster must not import core/http/auth. The comment now points at the spec that holds them together instead of claiming a derivation the move removed. - the epoch spec asserted e2 > e1, an ordering Claim's doc tells callers not to rely on. It asserts uniqueness, which is what the fence guarantees, and is named for that. A sibling spec still described the epoch as incrementing in SQL when it is drawn from a sequence. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/http/auth/public_routes.go | 12 +++-- core/services/cluster/instance.go | 59 ++++++++++++++++++------ core/services/cluster/instance_test.go | 18 ++++++-- core/services/cluster/membership.go | 19 ++++++-- core/services/cluster/membership_test.go | 25 ++++++++++ core/services/cluster/ownership_test.go | 34 ++++++++++++-- core/services/cluster/splice.go | 18 ++++++-- core/services/cluster/splice_test.go | 3 ++ core/services/cluster/wsconn.go | 2 +- 9 files changed, 158 insertions(+), 32 deletions(-) diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index 04d90d07d56c..4a8fcff01cac 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -76,10 +76,14 @@ func isPublicRoute(method, path string) bool { // ClusterPathPrefix is the replica-to-replica namespace. Its handlers check the // cluster token in the Authorization header themselves, so the check below lets -// them through the global session middleware. The cluster endpoints build their -// route paths from this same constant, which is why it lives here beside the -// check rather than beside the handlers: the exemption and the route it exempts -// cannot then be changed independently. +// them through the global session middleware. +// +// The cluster routes do NOT derive their paths from this constant: they are +// registered from core/services/cluster's own literal, because that package +// must not import core/http/auth. Nothing in the compiler holds the two +// together, so a spec does instead, driving a peer request through this +// middleware in core/http/endpoints/cluster/peer_test.go. Moving either string +// without the other turns that spec red, which is the whole reason it exists. const ClusterPathPrefix = "/api/cluster/" // usesAlternativeAuthentication identifies requests whose credentials are diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index d7e7afa71740..4140ce6bc796 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -128,13 +128,16 @@ func (r *Registry) Get(ctx context.Context, id string) (*Instance, error) { // address to advertise. The caller supplies the port, since the frontend's // listening port has nothing to do with the database's. // -// The discovery only holds while the database is a shared, remote host. When -// PostgreSQL runs on this same host or pod (compose, single-node, any sidecar -// layout) the route to it is loopback, and advertising 127.0.0.1 would make a -// peer dialling this replica reach itself instead. So an unspecified, loopback, -// or link-local 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. +// 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. @@ -172,16 +175,29 @@ func unroutableReason(ip net.IP, zone string) string { 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) - // A zone is only ever attached to a scoped (link-local) address, so this is - // the same rejection stated twice; the Zone check keeps the guarantee if a - // platform ever hands back a scoped address of another class, because - // IP.String() would silently drop the %iface and yield an undialable host. - case ip.IsLinkLocalUnicast() || zone != "": - return fmt.Sprintf("link-local (%s), which peers on other hosts cannot dial", 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. // @@ -207,13 +223,28 @@ func CheckAdvertisedAddr(addr string) (reason string, err error) { 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, ""), 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 diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go index d27043e03c46..aa0349496161 100644 --- a/core/services/cluster/instance_test.go +++ b/core/services/cluster/instance_test.go @@ -99,9 +99,11 @@ var _ = Describe("Advertised address discovery", func() { Expect(err).To(HaveOccurred()) }) - // A database on this same host routes over loopback on every platform, so - // this is deterministic rather than host-dependent. Returning 127.0.0.1 - // would make a peer dialling this replica reach itself. + // 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()) @@ -158,4 +160,14 @@ var _ = Describe("Checking a configured advertised address", func() { 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/membership.go b/core/services/cluster/membership.go index f3f71a321acd..ab24acc73e83 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -183,16 +183,24 @@ func (m *Membership) tick(ctx context.Context) { // would point every reader at an owner that no longer exists. This is the // announced form of what the sweeper does by inference, and the two must not // disagree about what "gone" removes. +// +// Instances first, then connections, which is deliberate and is the same order +// ReapStale takes. The two paths run concurrently in the ordinary case, a +// replica shutting down while a peer is sweeping it, and each locks the same +// two tables; opposite orders would let each hold the row the other is waiting +// for. PostgreSQL breaks such a cycle by aborting one side, so the cost is a +// failed shutdown rather than lost data, but an inversion that costs nothing to +// remove should not be left in. func (r *Registry) Deregister(ctx context.Context, id string) error { if err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - if err := tx.Where("owner_instance_id = ?", id).Delete(&NodeConnection{}).Error; err != nil { - return fmt.Errorf("deleting connections owned by %q: %w", id, err) - } // No RowsAffected check: deregistering a row another replica already // swept is the normal outcome of a slow shutdown, not an error. if err := tx.Where("id = ?", id).Delete(&Instance{}).Error; err != nil { return fmt.Errorf("deleting instance %q: %w", id, err) } + if err := tx.Where("owner_instance_id = ?", id).Delete(&NodeConnection{}).Error; err != nil { + return fmt.Errorf("deleting connections owned by %q: %w", id, err) + } return nil }); err != nil { return fmt.Errorf("deregistering instance %q: %w", id, err) @@ -225,6 +233,11 @@ func (r *Registry) Deregister(ctx context.Context, id string) error { // PostgreSQL only, like Live: distributed mode requires it, and the interval // arithmetic is measured on the database's clock because liveness is compared // across replicas. +// +// Instances are deleted before connections, and Deregister takes the same order +// on purpose, so the two paths cannot deadlock against each other. Here the +// order is also forced: the connection delete asks which instance rows survived, +// so it has to run second. func (r *Registry) ReapStale(ctx context.Context, self string, within time.Duration) (instances int64, connections int64, err error) { err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { res := tx.Where("id <> ? AND last_seen <= now() - make_interval(secs => ?)", self, within.Seconds()). diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go index d490d529dd47..c331eafe23d2 100644 --- a/core/services/cluster/membership_test.go +++ b/core/services/cluster/membership_test.go @@ -130,6 +130,31 @@ var _ = Describe("Reaping dead replicas", func() { 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", "10.0.0.2:8080", "v1")).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.deleteOrder()).To(Equal([]string{"instances", "node_connections"})) + Expect(reapRec.deleteOrder()).To(Equal(deregRec.deleteOrder()), + "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", "10.0.0.2:8080", "v1")).To(Succeed()) Expect(reg.Deregister(ctx, "gone")).To(Succeed()) diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index 6d6fdf58bc7e..cd59f8056f6c 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -57,6 +57,29 @@ func (r *sqlRecorder) only() string { return r.statements[0] } +// deleteOrder returns the tables the recorded statements deleted from, 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. +func (r *sqlRecorder) deleteOrder() []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 { + idx := strings.Index(strings.ToLower(stmt), "delete from ") + if idx < 0 { + continue + } + fields := strings.Fields(stmt[idx+len("delete from "):]) + if len(fields) == 0 { + continue + } + tables = append(tables, strings.Trim(fields[0], `"`)) + } + return tables +} + var _ = Describe("Connection ownership", func() { var ( db *gorm.DB @@ -71,12 +94,17 @@ var _ = Describe("Connection ownership", func() { reg = cluster.NewRegistry(db) }) - It("increments the epoch on every claim", func() { + 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).To(BeNumerically(">", e1)) + Expect(e2).ToNot(Equal(e1)) }) It("reports the latest owner", func() { @@ -160,7 +188,7 @@ var _ = Describe("Connection ownership", func() { Expect(err).To(MatchError(cluster.ErrNoConnection)) }) - It("claims in one statement that increments in SQL and stamps on the database clock", func() { + 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})) diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go index 5df97e890da0..604bbfec93d4 100644 --- a/core/services/cluster/splice.go +++ b/core/services/cluster/splice.go @@ -45,10 +45,12 @@ func Splice(a, b io.ReadWriteCloser) error { // Wait for the second direction so no copy is still touching either stream // once Splice has returned. This is load-bearing: it assumes Close unblocks // a copy parked in Read or Write, and a stream where that is false hangs - // here rather than leaking a goroutine. Both callers satisfy it. net.Conn - // does, and so does go-yamux/v5, whose Close sets readErr and calls - // notifyWaiting to wake a parked Read while a parked Write returns - // ErrStreamClosed. + // here rather than leaking a goroutine. The two stream types this is built + // for satisfy it: net.Conn does, and so does go-yamux/v5, whose Close sets + // readErr and calls notifyWaiting to wake a parked Read while a parked + // Write returns ErrStreamClosed. Nothing outside this package's own specs + // calls Splice yet, so a phase 2 caller relaying over anything else has to + // check this property rather than assume it. <-errs if first != nil { @@ -80,6 +82,14 @@ func copyStream(dst io.Writer, src io.Reader) error { // the peer is gone, and yamux produces exactly that when a Write races its // session's shutdown (see isMuxSessionFailure). // +// A socket-level abort (ECONNRESET, EPIPE) is deliberately absent too, which +// makes the same underlying event, a peer aborting mid-stream, reach the caller +// as nil over a yamux tunnel and as an error over a raw socket. That asymmetry +// is intended: the yamux endings are the teardown Splice's own Close provokes, +// so this primitive is the only thing that can tell them from a fault, whereas +// whether an aborted request is routine or a failure is the relay's policy and +// only the relay knows which request was abandoned. +// // The mux checks run first, and that ordering is load-bearing: a dying yamux // session usually hands every live stream its own cause wrapped up // (session.go:330), and that cause is routinely a closed-socket error, so diff --git a/core/services/cluster/splice_test.go b/core/services/cluster/splice_test.go index 7bbe3246de05..7b423a142730 100644 --- a/core/services/cluster/splice_test.go +++ b/core/services/cluster/splice_test.go @@ -308,6 +308,9 @@ var _ = Describe("Splice", func() { Expect(err).ToNot(HaveOccurred()) var far *yamux.Stream Eventually(accepted, "10s").Should(Receive(&far)) + // Deadline so a stream that never carries the byte fails this spec + // instead of parking the suite until its own timeout. + Expect(far.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) _, err = far.Read(make([]byte, 1)) Expect(err).ToNot(HaveOccurred()) diff --git a/core/services/cluster/wsconn.go b/core/services/cluster/wsconn.go index ca932d741cc8..aa8346de6453 100644 --- a/core/services/cluster/wsconn.go +++ b/core/services/cluster/wsconn.go @@ -119,7 +119,7 @@ func (c *wsConn) SetDeadline(t time.Time) error { func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) } // SetWriteDeadline takes writeMu because gorilla stores the write deadline in a -// plain struct field (conn.go:787) and applies it when it next flushes, so +// plain struct field (conn.go:796) and applies it when it next flushes, so // setting it while a write is in flight is a data race, not merely a late bound. func (c *wsConn) SetWriteDeadline(t time.Time) error { c.writeMu.Lock() From 5e2938ebf0f54a91cd0568d5b80113630c8b620c Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 06:46:02 +0000 Subject: [PATCH 17/42] feat(cluster): resolve tunnel ownership against a live owner OwnerRow is a bare row read of node_connections. A connection row outlives the replica that wrote it: a replica that dies stops heartbeating, but its rows survive until a peer's sweep removes them, which is up to InstanceLiveness plus one InstanceHeartbeat later. For that whole window the table names a process that is gone. The next component phase 2 builds is the relaying dialer, and a dialer reading OwnerRow would relay into a corpse for roughly 35 seconds after every replica death, then report the worker as unreachable when it is in fact absent, which is the distinction the phase 1 end-to-end specs pinned. Owner is the resolving read: one statement joining instances, returning ErrNoConnection when the row is missing OR its owner is not live. Both cases are one answer on purpose, since both mean no replica here holds this tunnel; they differ only in which sweep has run. It is one statement, not a row read followed by an instance lookup, because between two statements the owner can die and the caller would act on an owner the second read would have rejected. OwnerRow stays, unjoined, for readers that need the row itself, and a spec holds the two apart: with an aged-out owner, OwnerRow still names it and Owner refuses, so neither can quietly become the other. The liveness predicate is now one string, instanceIsLive, shared by Live and by Owner's join. Two spellings of one fact drift, and this drift would show as a relay to a replica one query calls dead and another calls alive. It is table-qualified so it is unambiguous inside the join, and the cutoff stays on the database clock, so replica clock skew cannot widen or narrow the window. Both mutations were run. Dropping the liveness predicate from the join fails 3 specs, the aged-owner one among them. Replacing the database clock with a Go-side time.Now() fails 1: the aged-owner specs still pass, because the two clocks agree on one host, and only the recorded-SQL spec sees the literal timestamp. That is why that spec exists. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/instance.go | 20 ++++- core/services/cluster/ownership.go | 55 ++++++++++-- core/services/cluster/ownership_test.go | 114 ++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 10 deletions(-) diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index 4140ce6bc796..c991d0da8d17 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -91,13 +91,27 @@ func (r *Registry) Heartbeat(ctx context.Context, id string) error { return nil } +// instanceIsLive is the one predicate that decides whether a replica is still +// alive, and it takes the window in seconds as its single bind parameter. Every +// reader of that fact uses this string: Live to list the survivors, Owner to +// refuse an owner that is not among them. Two spellings of one fact drift, and +// the drift would show up as a relay to a replica one query calls dead and +// another calls alive. +// +// The column is table-qualified because Owner reads it across a join, where an +// unqualified last_seen would be ambiguous. Postgres folds the unquoted name to +// the same table gorm quotes, so the qualification costs Live nothing. +// +// The cutoff is computed by the database for the same reason Register stamps +// there: liveness is compared across replicas, so a reader's own clock must not +// decide whether another replica is alive. +const instanceIsLive = `instances.last_seen > now() - make_interval(secs => ?)` + // Live returns the instances whose LastSeen is newer than now-within. func (r *Registry) Live(ctx context.Context, within time.Duration) ([]Instance, error) { var out []Instance - // The cutoff is computed by the database for the same reason Register stamps - // there: a reader's clock must not decide whether another replica is alive. if err := r.db.WithContext(ctx). - Where("last_seen > now() - make_interval(secs => ?)", within.Seconds()). + Where(instanceIsLive, within.Seconds()). Order("id"). Find(&out).Error; err != nil { return nil, fmt.Errorf("listing live instances: %w", err) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index 131b3d0b976d..48d9f94c58c9 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -168,14 +168,12 @@ func (r *Registry) Claim(ctx context.Context, nodeID, ownerID string) (int64, er // It answers "what does the table say", NOT "who holds this tunnel". The owner // it names may be dead: a replica that dies stops heartbeating, and its rows // survive until another replica's sweep removes them, which is up to -// InstanceLiveness plus one InstanceHeartbeat later. Any caller that ACTS on -// the answer must join instances itself and treat a non-live owner as -// ErrNoConnection; relaying to the row without that check is relaying into a -// process that is gone. +// InstanceLiveness plus one InstanceHeartbeat later. // -// The name says row on purpose, so that the joining version can take the plain -// name when phase 2 introduces the first caller that needs it. Nothing in -// phase 1 reads this outside tests, which is why the join is not here yet. +// Callers that ACT on the answer want Owner, which joins instances and treats a +// non-live owner as ErrNoConnection. This one is for callers that need to see +// the row itself, such as a sweeper deciding what to clean up, or a spec +// proving the two reads differ. func (r *Registry) OwnerRow(ctx context.Context, nodeID string) (string, int64, error) { var conn NodeConnection err := r.db.WithContext(ctx).Where("node_id = ?", nodeID).First(&conn).Error @@ -188,6 +186,49 @@ func (r *Registry) OwnerRow(ctx context.Context, nodeID string) (string, int64, return conn.OwnerInstanceID, conn.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 OwnerRow 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 and a dead owner are one answer on purpose. Both mean "no +// replica here 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. +// +// 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) { + var conn NodeConnection + err := r.db.WithContext(ctx). + Model(&NodeConnection{}). + // Only the connection's own columns are selected: the join exists to + // filter, and SELECT * across it would hand gorm the instances columns + // to scan into a NodeConnection. + Select("node_connections.*"). + Joins("JOIN instances ON instances.id = node_connections.owner_instance_id AND "+instanceIsLive, InstanceLiveness.Seconds()). + Where("node_connections.node_id = ?", 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 +} + // Release drops the claim identified by ownerID and epoch. Both are in the // WHERE so a replica that has only just noticed its dead socket cannot delete // the claim a later reconnect established elsewhere: the row it is trying to diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index cd59f8056f6c..b7abbbd6a896 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -188,6 +188,120 @@ var _ = Describe("Connection ownership", func() { Expect(err).To(MatchError(cluster.ErrNoConnection)) }) + // Owner is the resolving read: it answers "who holds this tunnel and can be + // relayed to", where OwnerRow 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() { + // Aged far enough past InstanceLiveness that the exact window boundary + // is not what these specs are measuring. + agedOut := 10 * time.Minute + + // 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", "10.0.0.1:8080", "v1")).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("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", "10.0.0.1:8080", "v1")).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", "10.0.0.1:8080", "v1")).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 through OwnerRow, which is why the two reads are separate", func() { + Expect(reg.Register(ctx, "inst-a", "10.0.0.1:8080", "v1")).To(Succeed()) + claimed, err := reg.Claim(ctx, "w1", "inst-a") + Expect(err).ToNot(HaveOccurred()) + age("inst-a", agedOut) + + owner, epoch, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred(), "OwnerRow reads 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 and OwnerRow must not agree here, or one of them is redundant") + }) + + 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", "10.0.0.1:8080", "v1")).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})) From e4777915f9fcf087b93eb5282c2272689fe96416 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 07:09:18 +0000 Subject: [PATCH 18/42] fix(cluster): make the ownership comments say only what holds Review round 1 on the joined Owner read. The behaviour was accepted; three comments claimed more than the code delivered, one spec pinned less than its doc promised, and one pre-existing spec ranked epochs. instanceIsLive said every reader of instance liveness uses it, which was false: ReapStale spelled the complement by hand. The complement is now written as NOT (instanceIsLive), so "stale" is exactly "not live", including how each side treats a NULL last_seen, and the sentence is true. Inverting that predicate fails 3 reaper specs, so the routing is held. The Select("node_connections.*") in Owner was justified by a SELECT * hazard that cannot occur: with a join present and nothing selected, gorm expands the model's own columns table-qualified (callbacks.BuildQuerySQL), and the suite is green with the Select removed. It stays, because the projection should be a property of this query, and the comment now says that instead. Owner gained the dialect guard Claim has. now() and make_interval are PostgreSQL, so on the SQLite single-binary path it failed with "no such function: now", which reads as a missing migration; that regression already shipped once in phase 1. The refusal is deliberately not ErrNoConnection: a deployment with no cluster has no answer about ownership, and reporting absence would let a caller conclude the worker is not connected. A spec in the non-PostgreSQL block holds all three properties. The new specs aged rows by ten minutes, which any window between zero and ten minutes satisfies, so nothing tied Owner's window to the one the sweeper uses. They now age to just past InstanceLiveness, and a sibling ages to half of it and must still resolve. Widening the window tenfold fails 2 specs, narrowing it tenfold fails 1; before this both were silent. The concurrent-claim spec asserted the stored epoch was the highest handed out, and justified it with claims drawing their epoch after the row lock, which contradicts Claim's own doc: the insert path draws nextval while the tuple is built. It now asserts the stored epoch is one of the epochs handed out, and ranks nothing. OwnerRow's doc justified the function with a sweeper that does not call it. ReapStale deletes orphans with a set difference; the callers are this package's specs and one e2e assertion. It says that, and states plainly that a caller needing to know who owns a node in order to dial it wants Owner. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/instance.go | 8 ++-- core/services/cluster/membership.go | 6 ++- core/services/cluster/ownership.go | 31 ++++++++++---- core/services/cluster/ownership_test.go | 56 ++++++++++++++++++------- 4 files changed, 74 insertions(+), 27 deletions(-) diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go index c991d0da8d17..606337a381ff 100644 --- a/core/services/cluster/instance.go +++ b/core/services/cluster/instance.go @@ -93,10 +93,10 @@ func (r *Registry) Heartbeat(ctx context.Context, id string) error { // instanceIsLive is the one predicate that decides whether a replica is still // alive, and it takes the window in seconds as its single bind parameter. Every -// reader of that fact uses this string: Live to list the survivors, Owner to -// refuse an owner that is not among them. Two spellings of one fact drift, and -// the drift would show up as a relay to a replica one query calls dead and -// another calls alive. +// reader of that fact is written in terms of it: Live lists the rows it selects, +// Owner refuses an owner it rejects, and ReapStale deletes its negation. Two +// spellings of one fact drift, and the drift would show up as a relay to a +// replica one query calls dead and another calls alive. // // The column is table-qualified because Owner reads it across a join, where an // unqualified last_seen would be ambiguous. Postgres folds the unquoted name to diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index ab24acc73e83..825dcf6e1cdb 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -240,7 +240,11 @@ func (r *Registry) Deregister(ctx context.Context, id string) error { // so it has to run second. func (r *Registry) ReapStale(ctx context.Context, self string, within time.Duration) (instances int64, connections int64, err error) { err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - res := tx.Where("id <> ? AND last_seen <= now() - make_interval(secs => ?)", self, within.Seconds()). + // Negated rather than spelled as its own comparison: "stale" has to be + // exactly "not live", including how each treats a row whose last_seen + // is NULL, and a hand-written complement is a second definition that + // only looks like the first. + 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) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index 48d9f94c58c9..8a80adede74b 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -170,10 +170,16 @@ func (r *Registry) Claim(ctx context.Context, nodeID, ownerID string) (int64, er // survive until another replica's sweep removes them, which is up to // InstanceLiveness plus one InstanceHeartbeat later. // -// Callers that ACT on the answer want Owner, which joins instances and treats a -// non-live owner as ErrNoConnection. This one is for callers that need to see -// the row itself, such as a sweeper deciding what to clean up, or a spec -// proving the two reads differ. +// Anything that needs to know WHO owns a node in order to act on it, a dialer +// above all, wants Owner: it joins instances and treats a non-live owner as +// ErrNoConnection. Dialing what this function returns is dialing a process that +// may be gone. +// +// What is left for this one is observing the table as such, independently of +// liveness. Its callers today are this package's specs, including the one that +// holds the two reads apart, and the e2e cluster spec that watches ownership +// move between replicas. No production caller reads it, and the sweeper is not +// one: ReapStale deletes orphans with a set difference in SQL. func (r *Registry) OwnerRow(ctx context.Context, nodeID string) (string, int64, error) { var conn NodeConnection err := r.db.WithContext(ctx).Where("node_id = ?", nodeID).First(&conn).Error @@ -210,12 +216,23 @@ func (r *Registry) OwnerRow(ctx context.Context, nodeID string) (string, int64, // 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{}). - // Only the connection's own columns are selected: the join exists to - // filter, and SELECT * across it would hand gorm the instances columns - // to scan into a 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()). Where("node_connections.node_id = ?", nodeID). diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index b7abbbd6a896..cd422199c638 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -193,9 +193,16 @@ var _ = Describe("Connection ownership", func() { // 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() { - // Aged far enough past InstanceLiveness that the exact window boundary - // is not what these specs are measuring. - agedOut := 10 * time.Minute + // 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 @@ -216,6 +223,18 @@ var _ = Describe("Connection ownership", func() { 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", "10.0.0.1:8080", "v1")).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 @@ -350,21 +369,16 @@ var _ = Describe("Connection ownership", func() { } Expect(seen).To(HaveLen(claimants)) - // Exactly one row, and its epoch is the highest handed out. That last - // part holds here because no Release intervenes: every claim after the - // first blocks on the row lock and draws its epoch after taking it, in - // commit order. It is not a general guarantee about epochs, which are - // unique but unordered. + // 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)) - var max int64 - for e := range seen { - if e > max { - max = e - } - } - Expect(rows[0].Epoch).To(Equal(max)) + Expect(seen).To(HaveKey(rows[0].Epoch), "the stored epoch was never handed to any claimant") }) }) @@ -387,6 +401,18 @@ var _ = Describe("Connection ownership on a non-PostgreSQL dialect", func() { 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()) From f3ba1f692b8e840fd21cfb463aa2d16247a35241 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 07:39:49 +0000 Subject: [PATCH 19/42] feat(cluster): hold worker tunnels, and re-claim them after a sweep Phase 1 left the connection fence with a table and no sockets behind it. This adds the registry that holds them: Attach claims the node and then stores the session, Open hands out a stream over the tunnel this replica holds, Detach releases the claim it was handed, and Held names what this process is carrying. The claim is written before the session is stored. A claimant that installs itself and only then finds it cannot claim has, for that window, published a tunnel no row records, so Held names it while a peer asking Owner is told the worker is connected nowhere. ErrNotOwner is produced at one place, the map miss. It is a routing fact: some other replica may hold that worker perfectly well. A broken socket under a held entry is returned as itself, because answering "not held here" would send a dialer looking elsewhere for a worker this replica is holding. Epochs are compared for equality and never ordered. Claim guarantees an epoch is unique and never reissued; it does not guarantee the later claim draws the larger number, because the sequence value on the insert path is drawn before the row lock. The membership loop now re-claims on re-register, which closes the hole phase 1 named in ReapStale. A replica that stalls long enough is swept by a peer, losing its instance row and, in the same transaction, every connection it owned; Register rebuilds the instance row and nothing else, so without this it serves workers that every other replica reports as connected nowhere. Re-claiming draws a fresh epoch, so an attachment carries two: the token Attach handed back, which is what Detach matches and which never moves, and the epoch of the row currently held, which is what Release is given. Collapsing them would leave the re-claimed row outliving the socket with no caller able to remove it. A tunnel whose session is already closed is skipped rather than claimed back, because claiming is an upsert and would take the row from whoever holds the worker now. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/membership.go | 68 ++++- core/services/cluster/tunnel.go | 278 ++++++++++++++++++ core/services/cluster/tunnel_test.go | 414 +++++++++++++++++++++++++++ 3 files changed, 749 insertions(+), 11 deletions(-) create mode 100644 core/services/cluster/tunnel.go create mode 100644 core/services/cluster/tunnel_test.go diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index 825dcf6e1cdb..ff511cdd9ffb 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -50,9 +50,11 @@ type Membership struct { done chan struct{} stopOnce sync.Once - // mu guards started, which tells Stop whether there is a loop to join. + // mu guards started, which tells Stop whether there is a loop to join, and + // tunnels, which SetTunnels may write while the loop is already reading it. mu sync.Mutex started bool + tunnels *TunnelRegistry } // NewMembership returns the membership loop for one replica. The address is @@ -75,6 +77,20 @@ func NewMembership(reg *Registry, id, addr, version string) *Membership { // registration is synchronous and its failure is returned: a replica whose // address never reaches the table is invisible to its peers, and starting // anyway would hide that behind a background log line. +// SetTunnels gives the loop the registry holding this replica's worker tunnels, +// so it can re-claim them after its rows have been swept. A Membership without +// one still heartbeats and sweeps; it simply has nothing to re-claim, which is +// the single-binary case and the case of a replica that accepts no tunnels. +// +// It is a setter rather than a constructor argument because the tunnel registry +// is what the tunnel endpoint is built on, and that is wired after membership +// is already running. +func (m *Membership) SetTunnels(t *TunnelRegistry) { + m.mu.Lock() + defer m.mu.Unlock() + m.tunnels = t +} + func (m *Membership) Start(ctx context.Context) error { if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { return err @@ -155,13 +171,20 @@ func (m *Membership) tick(ctx context.Context) { // enough to look dead. Re-register rather than heartbeat: a heartbeat // carries no address, so the row has to be rebuilt from scratch. // - // This rebuilds the instance row ONLY. The sweep that removed it also - // removed the connections this replica owned, and re-claiming those - // needs the tunnel registry phase 2 introduces; see ReapStale. + // Register rebuilds the instance row ONLY. The sweep that removed it + // removed the connections this replica owned in the same transaction, + // so the tunnels still held here have to be claimed again or this + // replica serves workers that, as far as every other replica can see, + // are connected nowhere. xlog.Warn("Cluster instance row was reaped, re-registering", "id", m.id) if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { xlog.Error("Re-registering cluster instance failed", "id", m.id, "error", err) + // Nothing to re-claim onto: a claim naming an instance row that + // does not exist is deleted by the next sweep that runs, this + // replica's own included, and the sweep is what has just happened. + return } + m.reclaimTunnels(ctx) } else if err != nil { xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) } @@ -176,6 +199,30 @@ func (m *Membership) tick(ctx context.Context) { } } +// reclaimTunnels re-writes a claim for every worker tunnel this replica still +// holds, after the sweep that deleted them. It is separate from tick only so +// the lock around the registry reference is not held across the database work. +func (m *Membership) reclaimTunnels(ctx context.Context) { + m.mu.Lock() + tunnels := m.tunnels + m.mu.Unlock() + if tunnels == nil { + return + } + + reclaimed, err := tunnels.Reclaim(ctx) + if err != nil { + // Logged rather than returned, and the loop keeps running: the next + // heartbeat fails the same way if the row is still missing, so the + // re-claim is retried. A worker whose claim never lands is reachable + // only through the replica it is connected to, which is this one. + xlog.Error("Re-claiming worker tunnels after this replica was reaped failed", "id", m.id, "error", err) + } + if reclaimed > 0 { + xlog.Info("Re-claimed worker tunnels after this replica was reaped", "id", m.id, "tunnels", reclaimed) + } +} + // Deregister removes one replica and the connections it owned. // // It deletes both, in one transaction, for the same reason ReapStale does: a @@ -222,13 +269,12 @@ func (r *Registry) Deregister(ctx context.Context, id string) error { // row would then delete the connections of workers that are, at that moment, // connected to it. // -// That protection is one-sided, and only the instance row recovers on its own. -// A replica that stalls long enough is reaped BY ANOTHER replica, taking its -// connection rows with it, and the re-register in tick rebuilds the instance -// row and nothing else: the sockets are still held here while the table says -// nobody holds them. Phase 2 closes this by re-claiming, on re-register, every -// connection this replica still holds locally, which needs the tunnel registry -// that owns those sockets. +// That protection is one-sided. A replica that stalls long enough is reaped BY +// ANOTHER replica, taking its connection rows with it, and Register rebuilds +// the instance row and nothing else. What restores the rest is the re-claim in +// tick, which writes a fresh claim for every tunnel the tunnel registry still +// holds; until it runs, this replica holds sockets the table records nobody +// holding. // // PostgreSQL only, like Live: distributed mode requires it, and the interval // arithmetic is measured on the database's clock because liveness is compared diff --git a/core/services/cluster/tunnel.go b/core/services/cluster/tunnel.go new file mode 100644 index 000000000000..7ca928a5e7b9 --- /dev/null +++ b/core/services/cluster/tunnel.go @@ -0,0 +1,278 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "context" + "errors" + "fmt" + "net" + "sort" + "sync" + "time" + + "github.com/libp2p/go-yamux/v5" + "github.com/mudler/xlog" +) + +// ErrNotOwner reports that this replica does not hold the tunnel for a node. +// +// It is a ROUTING fact and nothing else: some other replica may hold that +// worker perfectly well, and a caller that sees it relays through the owner the +// database names. It must therefore never be produced by anything that merely +// failed. A database error, a broken socket, a session that shut down under a +// held entry: each of those is reported as itself, because reporting them as +// "not held here" tells a dialer to look elsewhere for a worker that is right +// here, and the design forbids absence standing in for unreachable. This is the +// same rule Claim and Owner follow when they refuse a dialect rather than +// answering ErrNoConnection. +var ErrNotOwner = errors.New("cluster: this replica does not hold the tunnel for that node") + +// tunnelReleaseTimeout bounds the release Detach performs. Detach is called +// from the goroutine that has just watched a worker's session die, and that +// goroutine must not be parked on a database that went away with it. +const tunnelReleaseTimeout = 5 * time.Second + +// TunnelRegistry holds the worker tunnels this replica has accepted, and keeps +// the node_connections table agreeing with what it holds. +// +// It is the local half of the connection fence: the table says which replica +// owns a worker, and this says which socket that ownership actually resolves +// to. The two are written in one order, always, by Attach. +type TunnelRegistry struct { + reg *Registry + selfID string + + mu sync.Mutex + tunnels map[string]*heldTunnel +} + +// heldTunnel is one accepted worker tunnel. +// +// The two epochs are the same number until this replica is swept and re-claims, +// and they are separate fields because they answer different questions. +// +// token is what Attach handed back, and it is the only value Detach matches. +// It identifies one local attachment for that attachment's whole life, which is +// what lets a superseded holder's Detach be recognised as stale: epochs are +// never reissued, so a token from an earlier attachment cannot collide with a +// later one's. +// +// claim is the epoch of the row this replica currently holds for the node, and +// it is what Release must be given, because that is the row the fence matches +// on. A re-claim draws a fresh epoch and moves this one; leaving Release to use +// the token instead would match nothing, and the row would outlive the socket +// with no caller able to tell. +// +// Neither is ever ordered against the other, or against anything else. Claim +// guarantees uniqueness, not monotonicity. +type heldTunnel struct { + sess *yamux.Session + token int64 + claim int64 +} + +// NewTunnelRegistry returns a registry that claims tunnels as selfID. The ID +// must be the same one this replica registers in the instances table, since +// that is what Owner joins a claim against to decide the owner is alive. +func NewTunnelRegistry(reg *Registry, selfID string) *TunnelRegistry { + return &TunnelRegistry{ + reg: reg, + selfID: selfID, + tunnels: map[string]*heldTunnel{}, + } +} + +// Attach records this replica as the owner of nodeID's tunnel and stores the +// session, returning the epoch the caller must later hand to Detach. +// +// The claim is written BEFORE the session is stored, and the order is the +// point. A claimant that installs itself first and only then finds it cannot +// claim has, for that window, published a tunnel no row records: Held names it, +// and a peer asking Owner is told the worker is connected nowhere. Claiming +// first means a failed claim leaves this replica exactly as it was. +// +// A worker that re-dials onto this same replica supersedes its own earlier +// attachment, and the superseded session is closed here. Nothing else can close +// it: whoever accepted it is parked in AcceptStream on a session that is not +// broken, only replaced, and would wait there until the far side noticed. The +// caller keeps ownership of the session it passed in; only a session this +// registry evicted is closed by this registry. +func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux.Session) (int64, error) { + if sess == nil { + // Claiming would publish a tunnel that cannot carry anything, and the + // fence would then have to be unwound by a Detach nobody will call. + return 0, fmt.Errorf("attaching tunnel for node %q: no session", 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{sess: sess, token: epoch, claim: epoch} + t.mu.Unlock() + + if previous != nil && previous.sess != sess { + xlog.Debug("worker re-dialled this replica, dropping its previous tunnel", "node", nodeID) + _ = previous.sess.Close() + } + return epoch, nil +} + +// Detach drops the attachment epoch identifies and releases its claim. An epoch +// that is not the one Attach handed the current holder is a no-op, which is how +// a superseded holder noticing its dead socket is stopped from evicting the +// attachment that replaced it. +// +// Matched by EQUALITY, never by order. An epoch is unique and never reissued, +// but a claim inserted after a Release can draw a lower number than one already +// issued, so a stale token may compare either way against the live one. +// +// Releasing a claim this replica no longer holds is ordinary rather than +// exceptional: it is what a worker having re-homed to another replica looks +// like from here, so it is logged and not returned. Detach has no error to +// return to, being the last thing a dying tunnel's goroutine does. +func (t *TunnelRegistry) Detach(nodeID string, epoch int64) { + t.mu.Lock() + held, ok := t.tunnels[nodeID] + if !ok || held.token != epoch { + t.mu.Unlock() + return + } + delete(t.tunnels, nodeID) + claim := held.claim + t.mu.Unlock() + + // Not the caller's context, and not the one Attach was given: both belong + // to the request or the process that set the tunnel up, and by the time a + // tunnel is being torn down either may already be cancelled, which would + // leave the row behind on every ordinary disconnect. + ctx, cancel := context.WithTimeout(context.Background(), tunnelReleaseTimeout) + defer cancel() + // The claim, not the token: the row carries whatever epoch was last claimed + // for this attachment, and Release matches the row exactly. + if err := t.reg.Release(ctx, nodeID, t.selfID, claim); err != nil { + if errors.Is(err, ErrNoConnection) { + xlog.Debug("worker tunnel claim was already superseded", "node", nodeID, "epoch", claim) + return + } + xlog.Warn("Releasing a worker tunnel claim failed; peers will drop it when this replica's heartbeat ages out", + "node", nodeID, "epoch", claim, "error", err) + } +} + +// Open returns a stream to the worker over the tunnel this replica holds. +// +// ErrNotOwner means only that no tunnel for nodeID is held here. Every other +// failure is returned as itself, wrapped: a session that died under a held +// entry is a transport condition, and answering ErrNotOwner for it would send a +// dialer looking elsewhere for a worker this replica is holding. +// +// A failed open does not evict the entry. Whether a tunnel is held here is +// decided by Attach and Detach, and letting one bad open unhold it would race +// the goroutine that owns the session and is about to detach it properly. +func (t *TunnelRegistry) Open(ctx context.Context, nodeID string) (net.Conn, error) { + t.mu.Lock() + held, ok := t.tunnels[nodeID] + t.mu.Unlock() + if !ok { + return nil, fmt.Errorf("opening a stream to node %q: %w", nodeID, ErrNotOwner) + } + + stream, err := held.sess.OpenStream(ctx) + if err != nil { + return nil, fmt.Errorf("opening a stream to node %q over the tunnel held here: %w", nodeID, err) + } + return stream, nil +} + +// Held returns the nodes whose tunnels this replica holds, sorted. +// +// It answers what this process holds, which is not the same question as who the +// table says owns a node; Owner answers that one. The membership loop uses this +// to know what to re-claim after its rows have been swept. +func (t *TunnelRegistry) 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 fresh claim for every tunnel still held here, and returns +// how many it wrote. +// +// It exists 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 just re-registered. Re-registration rebuilds the instance row only, so +// without this the sockets are still held here while the table records nobody +// holding them, and every other replica answers "not connected" for workers +// that are connected. +// +// A closed session is skipped rather than claimed. Claiming is an upsert, so it +// takes the row from whoever holds it now, and a worker whose socket here is +// closed has already reconnected somewhere: claiming it back would point every +// dialer at a replica that cannot carry a byte to it. The check narrows that +// window rather than closing it, since a socket can be dead without this side +// having noticed; yamux's keepalive bounds how long that lasts, and the +// worker's next reconnect supersedes the claim in any case. +// +// The entry of a skipped tunnel is left alone. Whoever attached it owns its +// lifetime and will detach it; Reclaim is not an eviction path, and evicting +// here would race that goroutine. +// +// A single node's failure does not abort the rest: the tunnels are independent, +// and a claim that failed is retried on the next sweep this replica survives. +func (t *TunnelRegistry) Reclaim(ctx context.Context) (int, error) { + t.mu.Lock() + held := make(map[string]*heldTunnel, len(t.tunnels)) + for nodeID, tunnel := range t.tunnels { + held[nodeID] = tunnel + } + t.mu.Unlock() + + var reclaimed int + var errs []error + for nodeID, tunnel := range held { + if tunnel.sess.IsClosed() { + xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID) + continue + } + epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) + if err != nil { + errs = append(errs, err) + continue + } + t.mu.Lock() + // Re-read under the lock: the tunnel may have been detached, or + // superseded by a reconnect, while the claim was in flight. Writing the + // new claim onto whatever is there now would give a different + // attachment an epoch it never drew, and its Release would then match + // nothing. + // + // The epoch just drawn is then held by no attachment. If it also + // happened to be the last write to the row, the row outlives the + // attachment that no longer matches it: it still names this replica, + // truthfully, but no Detach will remove it. It is cleared by the + // worker's next claim anywhere, since Claim upserts, and otherwise by + // the sweep that eventually removes this replica. Serialising Attach + // and Reclaim per node would close it; the window is one database + // round trip during the tick that follows a sweep of this replica, and + // the machinery costs more than it removes. + if current, ok := t.tunnels[nodeID]; ok && current == tunnel { + current.claim = epoch + reclaimed++ + } + t.mu.Unlock() + } + if len(errs) > 0 { + return reclaimed, fmt.Errorf("re-claiming worker tunnels: %w", errors.Join(errs...)) + } + return reclaimed, nil +} diff --git a/core/services/cluster/tunnel_test.go b/core/services/cluster/tunnel_test.go new file mode 100644 index 000000000000..33f4c2d858c4 --- /dev/null +++ b/core/services/cluster/tunnel_test.go @@ -0,0 +1,414 @@ +package cluster_test + +import ( + "context" + "fmt" + "path/filepath" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +// workerTunnel returns the two halves of a worker's tunnel: the frontend holds +// the server half, because the worker is the side that dials. yamuxPair already +// builds exactly that pairing for peer links; this names the halves the way the +// worker path uses them so a spec cannot silently attach the wrong end. +func workerTunnel() (frontend *yamux.Session, worker *yamux.Session) { + GinkgoHelper() + worker, frontend = yamuxPair() + return frontend, worker +} + +// echoOnce accepts one stream on the worker's half and echoes what it reads. +// It is how a spec proves Open produced a stream that carries bytes, rather +// than a handle that merely exists. +func echoOnce(worker *yamux.Session) { + go func() { + defer GinkgoRecover() + stream, err := worker.AcceptStream() + if err != nil { + return + } + defer func() { _ = stream.Close() }() + buf := make([]byte, 4) + if _, err := stream.Read(buf); err != nil { + return + } + _, _ = stream.Write(buf) + }() +} + +// drain accepts and discards every stream on the worker's half, so a spec that +// opens streams without reading them does not park on the accept backlog. +func drain(worker *yamux.Session) { + go func() { + defer GinkgoRecover() + for { + stream, err := worker.AcceptStream() + if err != nil { + return + } + _ = stream.Close() + } + }() +} + +var _ = Describe("The worker tunnel registry", func() { + var ( + db *gorm.DB + reg *cluster.Registry + tun *cluster.TunnelRegistry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + tun = cluster.NewTunnelRegistry(reg, "me") + }) + + It("claims the node in the database before it stores the session", func() { + // A claimant that installs itself and only then tries to claim has, + // for that window, made the registry disagree with the table: Held + // names a tunnel no row records, and a peer asking Owner is told the + // worker is connected nowhere. The failure is injected through the + // production refusal path, a dialect with no epoch sequence, so the + // spec exercises the real error return rather than a fake. + sqliteDB, err := gorm.Open(sqlite.Open(filepath.Join(GinkgoT().TempDir(), "cluster.db")), &gorm.Config{}) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.Migrate(ctx, sqliteDB)).To(Succeed()) + unclaimable := cluster.NewTunnelRegistry(cluster.NewRegistry(sqliteDB), "me") + + frontend, _ := workerTunnel() + _, err = unclaimable.Attach(ctx, "w1", frontend) + Expect(err).To(HaveOccurred()) + + Expect(unclaimable.Held()).To(BeEmpty(), + "a claimant whose claim failed installed itself anyway, so the registry now disagrees with the table") + _, err = unclaimable.Open(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNotOwner)) + }) + + It("records the claim, so another replica can find the owner", func() { + frontend, _ := workerTunnel() + epoch, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + Expect(tun.Held()).To(ConsistOf("w1")) + owner, stored, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me")) + Expect(stored).To(Equal(epoch), "Attach handed back an epoch that is not the one it wrote") + }) + + It("opens a stream that carries bytes to the worker", func() { + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + echoOnce(worker) + + conn, err := tun.Open(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(conn).ToNot(BeNil()) + DeferCleanup(func() { _ = conn.Close() }) + + Expect(conn.SetDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = conn.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + buf := make([]byte, 4) + _, err = conn.Read(buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("ping")) + }) + + It("reports ErrNotOwner for a node whose tunnel it does not hold", func() { + // This is a routing fact and nothing more: some other replica may hold + // the worker perfectly well. It must never be produced by anything that + // merely failed. + _, err := tun.Open(ctx, "nobody") + Expect(err).To(MatchError(cluster.ErrNotOwner)) + }) + + It("supersedes an earlier attachment, and the superseded holder's Detach is a no-op", func() { + first, _ := workerTunnel() + firstEpoch, err := tun.Attach(ctx, "w1", first) + Expect(err).ToNot(HaveOccurred()) + + second, secondWorker := workerTunnel() + secondEpoch, err := tun.Attach(ctx, "w1", second) + Expect(err).ToNot(HaveOccurred()) + // Compared for difference, never for order. Claim guarantees an epoch + // is unique and never reissued; it does NOT guarantee the later claim + // draws the larger number, because the sequence value on the insert + // path is drawn before the row lock. + Expect(secondEpoch).ToNot(Equal(firstEpoch)) + + Expect(first.IsClosed()).To(BeTrue(), + "the superseded session was left open, so whoever is accepting on it never learns it was replaced") + + tun.Detach("w1", firstEpoch) + + Expect(tun.Held()).To(ConsistOf("w1"), "a stale Detach evicted the live session") + owner, stored, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me")) + Expect(stored).To(Equal(secondEpoch), "a stale Detach released the live claim") + + echoOnce(secondWorker) + conn, err := tun.Open(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Expect(conn.SetDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = conn.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + buf := make([]byte, 4) + _, err = conn.Read(buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("ping")) + + tun.Detach("w1", secondEpoch) + Expect(tun.Held()).To(BeEmpty()) + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("ignores a Detach naming an epoch it was never handed", func() { + frontend, _ := workerTunnel() + epoch, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + // Not an ordering probe: both directions are tried because an epoch is + // unique but unordered, so a stale token can compare either way against + // the live one and neither may be allowed to evict it. + tun.Detach("w1", epoch+1) + tun.Detach("w1", epoch-1) + + Expect(tun.Held()).To(ConsistOf("w1")) + _, stored, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(stored).To(Equal(epoch)) + }) + + It("does not report a held tunnel whose session has died as ErrNotOwner", func() { + // Absence and unreachability are different answers and callers act + // differently on them: a dialer told the worker is not here relays + // elsewhere or reports it gone, where the truth is that this replica + // holds the tunnel and the socket underneath it broke. + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + Expect(worker.Close()).To(Succeed()) + Eventually(frontend.IsClosed, "10s").Should(BeTrue()) + + _, err = tun.Open(ctx, "w1") + Expect(err).To(HaveOccurred()) + Expect(err).ToNot(MatchError(cluster.ErrNotOwner), + "a broken socket was reported as this replica not holding the tunnel") + Expect(tun.Held()).To(ConsistOf("w1"), + "holding the tunnel is a routing fact, and a failed Open is not what un-holds it") + }) + + It("serves Attach, Open, Held and Detach from independent goroutines", func() { + // Run under -race. The point is contention on one node's entry, not a + // tidy per-goroutine partition: a registry whose map is only ever + // touched by one goroutine at a time proves nothing about the one that + // is not. + const workers = 8 + start := make(chan struct{}) + var wg sync.WaitGroup + + attached := make(chan int64, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func(i int) { + defer GinkgoRecover() + defer wg.Done() + frontend, worker := workerTunnel() + drain(worker) + <-start + epoch, err := tun.Attach(ctx, fmt.Sprintf("w%d", i%2), frontend) + Expect(err).ToNot(HaveOccurred()) + attached <- epoch + if conn, err := tun.Open(ctx, fmt.Sprintf("w%d", i%2)); err == nil { + _ = conn.Close() + } + }(i) + } + readers := make(chan struct{}) + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer GinkgoRecover() + defer wg.Done() + <-start + for { + select { + case <-readers: + return + default: + tun.Held() + } + } + }() + } + + close(start) + epochs := make([]int64, 0, workers) + for i := 0; i < workers; i++ { + epochs = append(epochs, <-attached) + } + close(readers) + wg.Wait() + + // Exactly one attachment per node survived, whichever won, and the + // table agrees with the map about which. + Expect(tun.Held()).To(ConsistOf("w0", "w1")) + for _, node := range tun.Held() { + _, stored, err := reg.OwnerRow(ctx, node) + Expect(err).ToNot(HaveOccurred()) + Expect(epochs).To(ContainElement(stored), + "the table records an epoch no Attach ever handed out") + } + + // Every loser's Detach is a no-op; only the two winners empty the map. + for _, epoch := range epochs { + tun.Detach("w0", epoch) + tun.Detach("w1", epoch) + } + Expect(tun.Held()).To(BeEmpty()) + }) +}) + +var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", func() { + var ( + db *gorm.DB + reg *cluster.Registry + tun *cluster.TunnelRegistry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + tun = cluster.NewTunnelRegistry(reg, "me") + }) + + // reapSelf performs the two deletes a peer's sweep performs on this + // replica: the instance row and, in the same transaction, every connection + // it owned. Deregister is that transaction; ReapStale reaches it by aging + // last_seen, which cannot be done deterministically against a heartbeat + // loop that is refreshing the same column. That ReapStale deletes both is + // pinned separately, with no loop running, in the reaping specs. + reapSelf := func() { + GinkgoHelper() + Expect(reg.Deregister(ctx, "me")).To(Succeed()) + } + + It("re-claims every held tunnel when the loop finds its instance row gone", func() { + // Without this a replica that stalled long enough to be swept sits + // holding live worker sockets that no row records. Every other replica + // then answers "not connected" for workers that are connected, which is + // the absence-versus-unreachable failure the design forbids. + frontend, _ := workerTunnel() + epoch, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + membership := cluster.NewMembership(reg, "me", "10.0.0.1:8080", "v1") + membership.SetTunnels(tun) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + reapSelf() + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + + owner := func() (string, error) { + owner, _, err := reg.OwnerRow(ctx, "w1") + return owner, err + } + Eventually(owner, 3*cluster.InstanceHeartbeat, time.Second).Should(Equal("me")) + + _, reclaimed, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(reclaimed).ToNot(Equal(epoch), "the re-claim reused the epoch of a claim the sweep deleted") + + // The holder still carries the epoch Attach handed it, and is the only + // thing that will ever release this row. If Detach matched only the + // epoch the re-claim drew, the row would outlive the socket and no + // caller could tell. + tun.Detach("w1", epoch) + Expect(tun.Held()).To(BeEmpty()) + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("does not re-claim a tunnel whose session has already died", func() { + // Re-claiming is an upsert, so it takes the row from whoever holds it + // now. A worker whose socket here is dead has already reconnected + // somewhere, and claiming it back would point every dialer at a replica + // that cannot carry a byte to it. + live, _ := workerTunnel() + liveEpoch, err := tun.Attach(ctx, "live", live) + Expect(err).ToNot(HaveOccurred()) + dead, deadWorker := workerTunnel() + _, err = tun.Attach(ctx, "dead", dead) + Expect(err).ToNot(HaveOccurred()) + Expect(deadWorker.Close()).To(Succeed()) + Eventually(dead.IsClosed, "10s").Should(BeTrue()) + + membership := cluster.NewMembership(reg, "me", "10.0.0.1:8080", "v1") + membership.SetTunnels(tun) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + reapSelf() + owner := func() (string, error) { + owner, _, err := reg.OwnerRow(ctx, "live") + return owner, err + } + Eventually(owner, 3*cluster.InstanceHeartbeat, time.Second).Should(Equal("me")) + + _, reclaimedLive, err := reg.OwnerRow(ctx, "live") + Expect(err).ToNot(HaveOccurred()) + Expect(reclaimedLive).ToNot(Equal(liveEpoch)) + + _, _, err = reg.OwnerRow(ctx, "dead") + Expect(err).To(MatchError(cluster.ErrNoConnection), + "a tunnel whose session is closed was claimed back from whoever holds the worker now") + }) + + It("keeps re-claiming out of the ordinary heartbeat, which has nothing to rebuild", func() { + // A claim per tick would draw a fresh epoch every five seconds for + // every worker on this replica, and every one of those writes is a + // chance to take a row a reconnect has just moved elsewhere. + frontend, _ := workerTunnel() + epoch, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + membership := cluster.NewMembership(reg, "me", "10.0.0.1:8080", "v1") + membership.SetTunnels(tun) + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + stored := func() (int64, error) { + _, stored, err := reg.OwnerRow(ctx, "w1") + return stored, err + } + Consistently(stored, 2*cluster.InstanceHeartbeat, time.Second).Should(Equal(epoch)) + }) +}) From 63ed55a929bd22fe19f03d45212a28763d3e4738 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 08:32:51 +0000 Subject: [PATCH 20/42] fix(cluster): make a claim and its record indivisible per node Two Attach calls for one node both claim, and PostgreSQL serialises the two upserts, but nothing ordered the two map writes against the two commits. The entry left installed could be the one whose claim lost the row, and its Detach then released an epoch the row does not carry, so the release matched nothing and the row survived the socket. Nothing swept that row. This replica is alive and heartbeating, so ReapStale leaves its rows alone, and no reconnect is coming for a worker that has gone. Owner kept naming this replica as the live owner of a tunnel it no longer held, and every dialer sent here was answered ErrNotOwner, which is the relay into a replica that cannot serve the request that this phase exists to prevent. Claims for one node now pass through a gate, so claim and record are indivisible. It is per node rather than one lock over the registry, the way PeerPool locks per peer: the claim is a database round trip, and a slow one for a single worker must not hold up Open for every other. Detach is not gated, because it takes no context and must never park behind an in-flight database call, and it changes no epoch. Reclaim takes the same gate, which makes its claim the newest one for that node, so it records the epoch on whatever attachment is installed rather than only on the one it listed. Refusing to record onto an attachment that replaced the listed one would leave that row with nothing able to release it. The interleave the gate does not cover is Detach, and a claim whose attachment detached while it was in flight is now released again rather than left behind. Also: restore Start's doc comment, which SetTunnels had swallowed; keep reaping other replicas when this one fails to rebuild its own row, rather than skipping the sweep along with the re-claim; scope the comment about an unnoticed dead socket to the keepalive of the session whoever accepted the tunnel built, since the worker session config does not exist yet; and pin the sortedness of Held, the nil-session refusal, and both re-claim interleaves with specs. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/membership.go | 22 +-- core/services/cluster/tunnel.go | 199 ++++++++++++++++++---- core/services/cluster/tunnel_test.go | 239 +++++++++++++++++++++++++++ 3 files changed, 414 insertions(+), 46 deletions(-) diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index ff511cdd9ffb..4a5f3e19ccab 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -73,10 +73,6 @@ func NewMembership(reg *Registry, id, addr, version string) *Membership { } } -// Start registers this replica and begins heartbeating and sweeping. The first -// registration is synchronous and its failure is returned: a replica whose -// address never reaches the table is invisible to its peers, and starting -// anyway would hide that behind a background log line. // SetTunnels gives the loop the registry holding this replica's worker tunnels, // so it can re-claim them after its rows have been swept. A Membership without // one still heartbeats and sweeps; it simply has nothing to re-claim, which is @@ -91,6 +87,10 @@ func (m *Membership) SetTunnels(t *TunnelRegistry) { m.tunnels = t } +// Start registers this replica and begins heartbeating and sweeping. The first +// registration is synchronous and its failure is returned: a replica whose +// address never reaches the table is invisible to its peers, and starting +// anyway would hide that behind a background log line. func (m *Membership) Start(ctx context.Context) error { if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { return err @@ -177,14 +177,16 @@ func (m *Membership) tick(ctx context.Context) { // replica serves workers that, as far as every other replica can see, // are connected nowhere. xlog.Warn("Cluster instance row was reaped, re-registering", "id", m.id) - if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { + if err := m.reg.Register(ctx, m.id, m.addr, m.version); err == nil { + m.reclaimTunnels(ctx) + } else { + // Re-claiming is skipped and only re-claiming: a claim written now + // would name an instance row that does not exist, and the very next + // sweep deletes it as an orphan. The sweep below still runs, because + // what it removes is other replicas, and this replica failing to + // rebuild its own row is no reason to stop reaping theirs. xlog.Error("Re-registering cluster instance failed", "id", m.id, "error", err) - // Nothing to re-claim onto: a claim naming an instance row that - // does not exist is deleted by the next sweep that runs, this - // replica's own included, and the sweep is what has just happened. - return } - m.reclaimTunnels(ctx) } else if err != nil { xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) } diff --git a/core/services/cluster/tunnel.go b/core/services/cluster/tunnel.go index 7ca928a5e7b9..891bf77b8238 100644 --- a/core/services/cluster/tunnel.go +++ b/core/services/cluster/tunnel.go @@ -45,6 +45,10 @@ type TunnelRegistry struct { mu sync.Mutex tunnels map[string]*heldTunnel + // claiming holds one gate per node that a claim is in flight for. It is + // what makes "claim, then record the epoch" indivisible per node; see + // enterClaim. + claiming map[string]chan struct{} } // heldTunnel is one accepted worker tunnel. @@ -77,12 +81,66 @@ type heldTunnel struct { // that is what Owner joins a claim against to decide the owner is alive. func NewTunnelRegistry(reg *Registry, selfID string) *TunnelRegistry { return &TunnelRegistry{ - reg: reg, - selfID: selfID, - tunnels: map[string]*heldTunnel{}, + reg: reg, + selfID: selfID, + tunnels: map[string]*heldTunnel{}, + claiming: map[string]chan struct{}{}, } } +// enterClaim takes the gate for nodeID, so that no two claims for one node are +// ever in flight at the same time. leaveClaim releases it. +// +// It exists because a claim and the record of that claim are two steps, and +// between them the database has already moved. Two Attach calls for one node +// both claim, and PostgreSQL serialises the two upserts, but nothing orders the +// two map writes against the two commits: the entry that ends up installed can +// carry the epoch of the claim that did NOT win the row. Its Detach then +// releases an epoch the row does not hold, the release matches nothing, and the +// row survives the socket. Nothing sweeps that, because the replica named on it +// is alive and heartbeating, so Owner keeps naming this replica as the owner of +// a tunnel it no longer holds and every dialer routed here gets ErrNotOwner. +// +// The gate is per node rather than one lock over the whole registry so that a +// slow claim for one worker does not hold up Open for any other, the same +// reason PeerPool locks per peer. Detach is deliberately NOT gated: it takes no +// context and must never park behind an in-flight database call. It does not +// need to be, because it changes no epoch; what it can interleave with is +// covered where that matters, in Reclaim. +// +// The entry is deleted rather than kept, so the map holds only the claims +// actually in flight and cannot grow with the number of workers ever seen. +func (t *TunnelRegistry) 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() + + // Re-checked in the loop rather than taken on waking: several waiters + // are released by one close, and only one of them may proceed. + select { + case <-gate: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// leaveClaim releases the gate enterClaim took. The channel is closed rather +// than sent on, so every waiter wakes rather than one. +func (t *TunnelRegistry) 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 nodeID's tunnel and stores the // session, returning the epoch the caller must later hand to Detach. // @@ -98,6 +156,10 @@ func NewTunnelRegistry(reg *Registry, selfID string) *TunnelRegistry { // broken, only replaced, and would wait there until the far side noticed. The // caller keeps ownership of the session it passed in; only a session this // registry evicted is closed by this registry. +// +// Two Attach calls for one node are serialised, claim and record together, so +// the entry that survives is always the one whose claim the row carries. See +// enterClaim for what an unserialised pair leaves behind. func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux.Session) (int64, error) { if sess == nil { // Claiming would publish a tunnel that cannot carry anything, and the @@ -105,6 +167,11 @@ func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux. return 0, fmt.Errorf("attaching tunnel for node %q: no session", nodeID) } + if err := t.enterClaim(ctx, nodeID); err != nil { + return 0, fmt.Errorf("attaching tunnel for node %q: %w", nodeID, err) + } + defer t.leaveClaim(nodeID) + epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) if err != nil { return 0, err @@ -220,8 +287,9 @@ func (t *TunnelRegistry) Held() []string { // closed has already reconnected somewhere: claiming it back would point every // dialer at a replica that cannot carry a byte to it. The check narrows that // window rather than closing it, since a socket can be dead without this side -// having noticed; yamux's keepalive bounds how long that lasts, and the -// worker's next reconnect supersedes the claim in any case. +// having noticed, and how long that lasts is decided by the keepalive on the +// session whoever accepted the tunnel built. The worker's next reconnect +// supersedes the claim in any case. // // The entry of a skipped tunnel is left alone. Whoever attached it owns its // lifetime and will detach it; Reclaim is not an eviction path, and evicting @@ -231,48 +299,107 @@ func (t *TunnelRegistry) Held() []string { // and a claim that failed is retried on the next sweep this replica survives. func (t *TunnelRegistry) Reclaim(ctx context.Context) (int, error) { t.mu.Lock() - held := make(map[string]*heldTunnel, len(t.tunnels)) - for nodeID, tunnel := range t.tunnels { - held[nodeID] = tunnel + 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, tunnel := range held { - if tunnel.sess.IsClosed() { - xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID) - continue - } - epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) - if err != nil { + for _, nodeID := range held { + if err := t.reclaimOne(ctx, nodeID); err != nil { + if errors.Is(err, errTunnelNotReclaimed) { + continue + } errs = append(errs, err) continue } - t.mu.Lock() - // Re-read under the lock: the tunnel may have been detached, or - // superseded by a reconnect, while the claim was in flight. Writing the - // new claim onto whatever is there now would give a different - // attachment an epoch it never drew, and its Release would then match - // nothing. - // - // The epoch just drawn is then held by no attachment. If it also - // happened to be the last write to the row, the row outlives the - // attachment that no longer matches it: it still names this replica, - // truthfully, but no Detach will remove it. It is cleared by the - // worker's next claim anywhere, since Claim upserts, and otherwise by - // the sweep that eventually removes this replica. Serialising Attach - // and Reclaim per node would close it; the window is one database - // round trip during the tick that follows a sweep of this replica, and - // the machinery costs more than it removes. - if current, ok := t.tunnels[nodeID]; ok && current == tunnel { - current.claim = epoch - reclaimed++ - } - t.mu.Unlock() + reclaimed++ } if len(errs) > 0 { return reclaimed, fmt.Errorf("re-claiming worker tunnels: %w", errors.Join(errs...)) } return reclaimed, nil } + +// errTunnelNotReclaimed reports that a node was passed over rather than failed: +// its session is closed, or the attachment went away while the claim was in +// flight. It never leaves this file. It exists so Reclaim's count stays honest +// without "skipped" having to look like an error to its caller. +var errTunnelNotReclaimed = errors.New("cluster: tunnel not re-claimed") + +// reclaimOne writes a fresh claim for one node and records it on whatever +// attachment is installed for that node. +// +// Whatever is installed when the gate is taken, not whatever Reclaim listed a +// moment earlier. A worker that re-dialled in between has an entry carrying the +// epoch of ITS claim, and the gate makes that claim strictly older than this +// one, so the row now holds this epoch and only this entry can release it. +// Refusing to record onto an attachment because it is not the one listed would +// leave that row with no attachment able to release it, which is the leak this +// whole function exists to prevent. +// +// If nothing is installed, the attachment detached while the claim was in +// flight. Detach is not gated, so this is reachable, and it is the one case +// where a claim is drawn that no attachment will ever release: the row would +// name this replica for a tunnel it does not hold, and Owner would send every +// dialer here to be told ErrNotOwner. The claim is therefore released again. +// Releasing it cannot take anyone else's row, because Release matches the epoch +// exactly and no epoch is ever reissued. +func (t *TunnelRegistry) reclaimOne(ctx context.Context, nodeID string) error { + // The gate is taken before the entry is even read, so that everything this + // function decides is decided about the attachment its claim will land on. + // Reading first and gating after would leave a window in which a re-dial + // replaces the entry, and the liveness this checked would be a property of + // a session it is no longer claiming for. + if err := t.enterClaim(ctx, nodeID); err != nil { + return fmt.Errorf("re-claiming node %q: %w", nodeID, err) + } + + t.mu.Lock() + tunnel, ok := t.tunnels[nodeID] + t.mu.Unlock() + if !ok { + // Detached between Reclaim listing the nodes and this gate. Nothing was + // claimed, so there is nothing to undo. + t.leaveClaim(nodeID) + return errTunnelNotReclaimed + } + if tunnel.sess.IsClosed() { + xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID) + t.leaveClaim(nodeID) + return errTunnelNotReclaimed + } + + epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) + if err != nil { + t.leaveClaim(nodeID) + return err + } + + t.mu.Lock() + current, installed := t.tunnels[nodeID] + if installed { + // current is necessarily the entry read above: the gate is still held, + // and Attach and reclaimOne are the only writers that install one. The + // identity is therefore not re-checked; the case that IS reachable is + // the entry being gone, because Detach is not gated. + current.claim = epoch + } + t.mu.Unlock() + t.leaveClaim(nodeID) + + if installed { + return nil + } + + // Released outside the gate: it is a second round trip, and holding the + // gate across it would park a worker re-dialling this node behind a + // cleanup. A re-dial that claims first simply makes this release match + // nothing, which is the same no-op it would have been. + if err := t.reg.Release(ctx, nodeID, t.selfID, epoch); err != nil && !errors.Is(err, ErrNoConnection) { + return fmt.Errorf("releasing a re-claim for detached node %q: %w", nodeID, err) + } + return errTunnelNotReclaimed +} diff --git a/core/services/cluster/tunnel_test.go b/core/services/cluster/tunnel_test.go index 33f4c2d858c4..ac7ff4ad8c37 100644 --- a/core/services/cluster/tunnel_test.go +++ b/core/services/cluster/tunnel_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "path/filepath" + "strings" "sync" "time" @@ -15,8 +16,81 @@ import ( . "github.com/onsi/gomega" "gorm.io/driver/sqlite" "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" ) +// claimHook runs an action once, from inside the database call that issued a +// matching statement, on that call's own goroutine. +// +// It is how a spec pins an interleaving instead of racing for one. gorm calls +// its logger's Trace after the statement has executed and before the Create or +// Delete that issued it returns (gorm@v1.31.1/callbacks.go:139-145), with the +// bind values interpolated into the SQL, so an action installed here runs at +// the one instant a claim has been written and not yet recorded. Racing two +// goroutines and hoping to land in that window is the flaky spec this replaces. +// +// It fires at most once: the action itself issues statements through the same +// session, and an unguarded hook would recurse. +type claimHook struct { + gormlogger.Interface + mu sync.Mutex + fired bool + match func(sql string) bool + action func(sql string) +} + +func newClaimHook(match func(sql string) bool) *claimHook { + return &claimHook{Interface: gormlogger.Default.LogMode(gormlogger.Silent), match: match} +} + +func (h *claimHook) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) { + sql, rows := fc() + h.mu.Lock() + fire := !h.fired && h.action != nil && h.match(sql) + if fire { + h.fired = true + } + h.mu.Unlock() + if fire { + h.action(sql) + } + h.Interface.Trace(ctx, begin, func() (string, int64) { return sql, rows }, err) +} + +// isClaimOf matches the upsert Claim issues for one node. +func isClaimOf(nodeIDs ...string) func(string) bool { + return func(sql string) bool { + if !strings.Contains(sql, "INSERT INTO \"node_connections\"") { + return false + } + for _, nodeID := range nodeIDs { + if strings.Contains(sql, "'"+nodeID+"'") { + return true + } + } + return false + } +} + +// claimedNode reports which of the named nodes a claim statement was for. +func claimedNode(sql string, nodeIDs ...string) string { + GinkgoHelper() + for _, nodeID := range nodeIDs { + if strings.Contains(sql, "'"+nodeID+"'") { + return nodeID + } + } + Fail("the claim statement named none of " + strings.Join(nodeIDs, ", ")) + return "" +} + +// serializationProbe is how long a spec watches for something that must not +// happen. It bounds an assertion about an ABSENT event, which is the only kind +// of wait a spec cannot replace with a channel: there is no event to receive. +// The thing it watches for takes one database round trip when the serialisation +// it guards is missing, so this is orders of magnitude longer than it needs. +const serializationProbe = 500 * time.Millisecond + // workerTunnel returns the two halves of a worker's tunnel: the frontend holds // the server half, because the worker is the side that dials. yamuxPair already // builds exactly that pairing for peer links; this names the halves the way the @@ -220,6 +294,85 @@ var _ = Describe("The worker tunnel registry", func() { "holding the tunnel is a routing fact, and a failed Open is not what un-holds it") }) + It("refuses a nil session rather than claiming a tunnel that cannot carry anything", func() { + // A claim written for a session that does not exist publishes a tunnel + // to every replica in the deployment, and the fence would then have to + // be unwound by a Detach nobody is going to call. + _, err := tun.Attach(ctx, "w1", nil) + Expect(err).To(HaveOccurred()) + + Expect(tun.Held()).To(BeEmpty()) + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), + "a node with no session was published to the deployment as connected here") + }) + + It("returns the nodes it holds in sorted order", func() { + // Attached out of order on purpose: sorted output is what makes a log + // line and a re-claim pass comparable between two runs, and a map + // range would only look sorted until it did not. + for _, nodeID := range []string{"w3", "w1", "w2"} { + frontend, _ := workerTunnel() + _, err := tun.Attach(ctx, nodeID, frontend) + Expect(err).ToNot(HaveOccurred()) + } + Expect(tun.Held()).To(Equal([]string{"w1", "w2", "w3"})) + }) + + It("serialises two Attach calls for one node, so the surviving entry holds the row's epoch", func() { + // Two claims for one node are serialised by PostgreSQL, but nothing + // orders the two map writes against the two commits. Unserialised, the + // entry left installed can carry the epoch of the claim that lost the + // row: its Detach then releases an epoch the row does not hold, the + // release matches nothing, and the row outlives the socket. Nothing + // sweeps that, because this replica is alive and heartbeating, so Owner + // keeps sending dialers here to be told ErrNotOwner. + hook := newClaimHook(isClaimOf("w1")) + hooked := cluster.NewTunnelRegistry( + cluster.NewRegistry(db.Session(&gorm.Session{Logger: hook})), "me") + + secondSession, _ := workerTunnel() + secondStarted := make(chan struct{}) + secondEpochs := make(chan int64, 1) + hook.action = func(string) { + // Launched from inside the first claim, so the second Attach is + // provably reaching for the same node while the first is between + // its claim and its store. Starting it before the call would leave + // which one claims first to the scheduler. + go func() { + defer GinkgoRecover() + close(secondStarted) + epoch, err := hooked.Attach(ctx, "w1", secondSession) + Expect(err).ToNot(HaveOccurred()) + secondEpochs <- epoch + }() + <-secondStarted + Consistently(secondEpochs, serializationProbe, 10*time.Millisecond).ShouldNot(Receive(), + "a second Attach for this node claimed AND recorded its epoch while the first was between its own claim and store") + } + + firstSession, _ := workerTunnel() + firstEpoch, err := hooked.Attach(ctx, "w1", firstSession) + Expect(err).ToNot(HaveOccurred()) + var secondEpoch int64 + Eventually(secondEpochs, "10s").Should(Receive(&secondEpoch)) + Expect(secondEpoch).ToNot(Equal(firstEpoch)) + + _, stored, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect([]int64{firstEpoch, secondEpoch}).To(ContainElement(stored)) + + // Whichever attachment survived, one of these two Detach calls is the + // live one and must take the row with it. If neither does, the row is + // carrying an epoch no attachment holds. + hooked.Detach("w1", firstEpoch) + hooked.Detach("w1", secondEpoch) + Expect(hooked.Held()).To(BeEmpty()) + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), + "the surviving attachment could not release its row, so the row outlived the socket") + }) + It("serves Attach, Open, Held and Detach from independent goroutines", func() { // Run under -race. The point is contention on one node's entry, not a // tidy per-goroutine partition: a registry whose map is only ever @@ -288,6 +441,11 @@ var _ = Describe("The worker tunnel registry", func() { tun.Detach("w1", epoch) } Expect(tun.Held()).To(BeEmpty()) + for _, node := range []string{"w0", "w1"} { + _, _, err := reg.OwnerRow(ctx, node) + Expect(err).To(MatchError(cluster.ErrNoConnection), + "node %s kept a row no attachment could release", node) + } }) }) @@ -392,6 +550,87 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu "a tunnel whose session is closed was claimed back from whoever holds the worker now") }) + It("records the re-claim on the attachment installed now, not the one it listed", func() { + // Reclaim lists the nodes it holds, then claims them one at a time. A + // worker that re-dials in between leaves an entry the list never saw. + // The claim is drawn under that node's gate, so it is the NEWEST claim + // for the node and the row carries it: recording it on the entry that + // is installed is the only thing that lets that entry release the row. + // Refusing to record it because the entry is not the one listed would + // leave the row behind when the socket dies. + hook := newClaimHook(isClaimOf("w1", "w2")) + hooked := cluster.NewTunnelRegistry( + cluster.NewRegistry(db.Session(&gorm.Session{Logger: hook})), "me") + + for _, nodeID := range []string{"w1", "w2"} { + frontend, _ := workerTunnel() + _, err := hooked.Attach(ctx, nodeID, frontend) + Expect(err).ToNot(HaveOccurred()) + } + + // The re-dial lands on whichever node this pass has not reached yet, + // so the spec does not depend on which one Reclaim takes first. + type redial struct { + node string + epoch int64 + } + redialled := make(chan redial, 1) + hook.action = func(sql string) { + node := "w2" + if claimedNode(sql, "w1", "w2") == "w2" { + node = "w1" + } + frontend, _ := workerTunnel() + epoch, err := hooked.Attach(ctx, node, frontend) + Expect(err).ToNot(HaveOccurred()) + redialled <- redial{node: node, epoch: epoch} + } + + count, err := hooked.Reclaim(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(count).To(Equal(2)) + + var latest redial + Expect(redialled).To(Receive(&latest)) + _, stored, err := reg.OwnerRow(ctx, latest.node) + Expect(err).ToNot(HaveOccurred()) + Expect(stored).ToNot(Equal(latest.epoch), "the re-claim never reached the node that re-dialled") + + // The attachment that re-dialled is the one holding the socket, so its + // Detach has to be the one that removes the row. + hooked.Detach(latest.node, latest.epoch) + _, _, err = reg.OwnerRow(ctx, latest.node) + Expect(err).To(MatchError(cluster.ErrNoConnection), + "the re-claim was recorded on nothing, so the attachment that holds the socket cannot release its row") + }) + + It("releases a re-claim whose attachment detached while the claim was in flight", func() { + // Detach is not gated against a re-claim, so this interleave is real: + // the claim commits, then the socket dies and Detach releases the epoch + // it was given, which the claim has already replaced. Left alone, the + // row names this replica for a tunnel it no longer holds, nothing + // sweeps it because this replica is alive, and Owner sends every dialer + // here to be told ErrNotOwner. + hook := newClaimHook(isClaimOf("w1")) + hooked := cluster.NewTunnelRegistry( + cluster.NewRegistry(db.Session(&gorm.Session{Logger: hook})), "me") + + frontend, _ := workerTunnel() + epoch, err := hooked.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + hook.action = func(string) { hooked.Detach("w1", epoch) } + + count, err := hooked.Reclaim(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(count).To(BeZero(), "a node that detached mid-claim was counted as re-claimed") + + Expect(hooked.Held()).To(BeEmpty()) + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), + "the re-claim left a row behind that no attachment holds and no sweep will remove") + }) + It("keeps re-claiming out of the ordinary heartbeat, which has nothing to rebuild", func() { // A claim per tick would draw a fresh epoch every five seconds for // every worker on this replica, and every one of those writes is a From 62476e553ea9acd79a3e2e8414028545ba4a97d8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 09:06:42 +0000 Subject: [PATCH 21/42] fix(cluster): keep a session close out of the per-node claim gate The gate is justified by being held for one claim round trip, and Attach held it across the close of the session it superseded. Closing a yamux session closes the underlying conn and then waits for both its send and recv loops to exit, and the send loop can be inside a write bounded only by ConnectionWriteTimeout, so that is a wait on other goroutines. It must not stand between a worker re-dialling this node and its claim. The gate is now released after the store and before the close, which also makes Attach match reclaimOne, where it has always been released explicitly on every path. This is safe because a superseded session is no longer reachable from the map by the time it is closed: the next re-dial replaces an entry that already names the new session. Pin the re-claim half of the gate too. A worker that re-dials between a re-claim's commit and its record leaves the row carrying the re-dial's epoch while the entry carries the re-claim's, so the attachment holding the socket releases an epoch the row does not have and the row outlives it, with nothing to sweep it while this replica is alive. Only Attach's half of the serialisation was asserted; keying the two apart left every spec green. Also take the test hook's action under the lock that guards whether it has fired. It was written from the spec's goroutine and read from whichever goroutine issued the statement, which is a race in the harness that pins the serialisation specs. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/tunnel.go | 15 +++++- core/services/cluster/tunnel_test.go | 79 +++++++++++++++++++++++++--- 2 files changed, 86 insertions(+), 8 deletions(-) diff --git a/core/services/cluster/tunnel.go b/core/services/cluster/tunnel.go index 891bf77b8238..36f7c91d3d4b 100644 --- a/core/services/cluster/tunnel.go +++ b/core/services/cluster/tunnel.go @@ -170,10 +170,10 @@ func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux. if err := t.enterClaim(ctx, nodeID); err != nil { return 0, fmt.Errorf("attaching tunnel for node %q: %w", nodeID, err) } - defer t.leaveClaim(nodeID) epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) if err != nil { + t.leaveClaim(nodeID) return 0, err } @@ -181,7 +181,20 @@ func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux. previous := t.tunnels[nodeID] t.tunnels[nodeID] = &heldTunnel{sess: sess, token: epoch, claim: epoch} t.mu.Unlock() + t.leaveClaim(nodeID) + // Closed after the gate is released, not under it. The gate is justified by + // being held for one claim round trip, and closing a session is not that: + // yamux closes the underlying conn and then waits for both its send and + // recv loops to exit (go-yamux/v5@v5.1.0/session.go:330-332), and the send + // loop can be inside a write bounded only by ConnectionWriteTimeout. That + // is a wait on other goroutines, and it must not stand between a worker + // re-dialling this node and its claim. + // + // Releasing first is safe because the superseded session is no longer + // reachable from the map: whoever re-dials next replaces an entry that + // already names the new session, and this close can only ever affect the + // one it just displaced. if previous != nil && previous.sess != sess { xlog.Debug("worker re-dialled this replica, dropping its previous tunnel", "node", nodeID) _ = previous.sess.Close() diff --git a/core/services/cluster/tunnel_test.go b/core/services/cluster/tunnel_test.go index ac7ff4ad8c37..67993130ee54 100644 --- a/core/services/cluster/tunnel_test.go +++ b/core/services/cluster/tunnel_test.go @@ -43,16 +43,26 @@ func newClaimHook(match func(sql string) bool) *claimHook { return &claimHook{Interface: gormlogger.Default.LogMode(gormlogger.Silent), match: match} } +// setAction installs what the hook runs, under the same lock that guards fired. +// The action is written from the spec's goroutine and read from whichever +// goroutine issues the statement, which need not be the same one. +func (h *claimHook) setAction(action func(sql string)) { + h.mu.Lock() + defer h.mu.Unlock() + h.action = action +} + func (h *claimHook) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) { sql, rows := fc() h.mu.Lock() - fire := !h.fired && h.action != nil && h.match(sql) + action := h.action + fire := !h.fired && action != nil && h.match(sql) if fire { h.fired = true } h.mu.Unlock() if fire { - h.action(sql) + action(sql) } h.Interface.Trace(ctx, begin, func() (string, int64) { return sql, rows }, err) } @@ -334,7 +344,7 @@ var _ = Describe("The worker tunnel registry", func() { secondSession, _ := workerTunnel() secondStarted := make(chan struct{}) secondEpochs := make(chan int64, 1) - hook.action = func(string) { + hook.setAction(func(string) { // Launched from inside the first claim, so the second Attach is // provably reaching for the same node while the first is between // its claim and its store. Starting it before the call would leave @@ -349,7 +359,7 @@ var _ = Describe("The worker tunnel registry", func() { <-secondStarted Consistently(secondEpochs, serializationProbe, 10*time.Millisecond).ShouldNot(Receive(), "a second Attach for this node claimed AND recorded its epoch while the first was between its own claim and store") - } + }) firstSession, _ := workerTunnel() firstEpoch, err := hooked.Attach(ctx, "w1", firstSession) @@ -575,7 +585,7 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu epoch int64 } redialled := make(chan redial, 1) - hook.action = func(sql string) { + hook.setAction(func(sql string) { node := "w2" if claimedNode(sql, "w1", "w2") == "w2" { node = "w1" @@ -584,7 +594,7 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu epoch, err := hooked.Attach(ctx, node, frontend) Expect(err).ToNot(HaveOccurred()) redialled <- redial{node: node, epoch: epoch} - } + }) count, err := hooked.Reclaim(ctx) Expect(err).ToNot(HaveOccurred()) @@ -604,6 +614,61 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu "the re-claim was recorded on nothing, so the attachment that holds the socket cannot release its row") }) + It("serialises a re-claim against an Attach for the same node", func() { + // The re-claim takes the same gate Attach does, and for the same + // reason. Unserialised, a worker re-dialling between the re-claim's + // commit and its record leaves the row carrying the re-dial's epoch + // while the entry carries the re-claim's, so the attachment holding the + // socket releases an epoch the row does not have. Nothing sweeps the + // row that is left, because this replica is alive: Owner keeps naming + // it as the owner of a tunnel it does not hold. + hook := newClaimHook(isClaimOf("w1")) + hooked := cluster.NewTunnelRegistry( + cluster.NewRegistry(db.Session(&gorm.Session{Logger: hook})), "me") + + first, _ := workerTunnel() + attachEpoch, err := hooked.Attach(ctx, "w1", first) + Expect(err).ToNot(HaveOccurred()) + + redialSession, _ := workerTunnel() + redialStarted := make(chan struct{}) + redialEpochs := make(chan int64, 1) + hook.setAction(func(string) { + // Launched from inside the re-claim's own claim, so the re-dial is + // provably reaching for this node while the re-claim is between its + // commit and its record. + go func() { + defer GinkgoRecover() + close(redialStarted) + epoch, err := hooked.Attach(ctx, "w1", redialSession) + Expect(err).ToNot(HaveOccurred()) + redialEpochs <- epoch + }() + <-redialStarted + Consistently(redialEpochs, serializationProbe, 10*time.Millisecond).ShouldNot(Receive(), + "a worker re-dialled and recorded its claim while a re-claim for the same node was between its own claim and record") + }) + + count, err := hooked.Reclaim(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(count).To(Equal(1)) + + var redialEpoch int64 + Eventually(redialEpochs, "10s").Should(Receive(&redialEpoch)) + Expect(redialEpoch).ToNot(Equal(attachEpoch)) + + // The re-dial claimed after the re-claim, so the row carries its epoch + // and its attachment is the one that has to be able to release it. The + // superseded token must still be a no-op. + hooked.Detach("w1", attachEpoch) + Expect(hooked.Held()).To(ConsistOf("w1")) + hooked.Detach("w1", redialEpoch) + Expect(hooked.Held()).To(BeEmpty()) + _, _, err = reg.OwnerRow(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection), + "the attachment that re-dialled could not release its row, so the row outlived the socket") + }) + It("releases a re-claim whose attachment detached while the claim was in flight", func() { // Detach is not gated against a re-claim, so this interleave is real: // the claim commits, then the socket dies and Detach releases the epoch @@ -619,7 +684,7 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu epoch, err := hooked.Attach(ctx, "w1", frontend) Expect(err).ToNot(HaveOccurred()) - hook.action = func(string) { hooked.Detach("w1", epoch) } + hook.setAction(func(string) { hooked.Detach("w1", epoch) }) count, err := hooked.Reclaim(ctx) Expect(err).ToNot(HaveOccurred()) From 6e55092a4b5d1c37e7ffa4d60974b40541cbe212 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 09:42:26 +0000 Subject: [PATCH 22/42] feat(cluster): open the door a worker dials its tunnel through A worker needs no inbound port: it dials GET /api/cluster/connect, the connection becomes one multiplexed yamux session, and the frontend opens a stream on it per request. This adds the endpoint that accepts that dial and attaches it to the tunnel registry. The dial is authenticated against the NODE's own stored token hash rather than the deployment's registration token. That is the mechanism, not yet the isolation, since a worker still registers by presenting the shared token; what it rules out is the shortcut of comparing against the configured value, which would have to be unpicked the day workers get their own secrets. Every refusal happens BEFORE the WebSocket upgrade, so a dialer reads an HTTP status rather than a handshake error. The route is registered in every deployment, single-binary ones included, which is what puts it in front of the route-coverage test that holds that rule in place; with no node registry it refuses every dial, and tells a credentialed one the frontend has no cluster rather than that its token is wrong. A lookup that FAILED is answered as a failure. Reporting a database that could not be read as "unauthorized" would send a worker re-registering, throwing away the identity its tunnel and loaded models are keyed by. Wires the tunnel registry in core/application/distributed.go and hands it to the membership loop. Without that call the re-claim after a replica is reaped had no production caller and could never run. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] --- core/application/distributed.go | 26 ++ core/http/app.go | 12 + core/http/endpoints/cluster/connect.go | 191 +++++++++++++ core/http/endpoints/cluster/connect_test.go | 285 ++++++++++++++++++++ core/http/routes/cluster.go | 20 ++ core/services/cluster/tunnel.go | 11 + docs/content/features/distributed-mode.md | 18 ++ 7 files changed, 563 insertions(+) create mode 100644 core/http/endpoints/cluster/connect.go create mode 100644 core/http/endpoints/cluster/connect_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 8769f006d228..30b297913f06 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -56,6 +56,11 @@ type DistributedServices struct { Membership *cluster.Membership // PeerSessions owns the peer links other replicas dialled into this one. PeerSessions *cluster.SessionStore + // Tunnels holds the worker tunnels this replica has accepted and keeps the + // node_connections table agreeing with them. It is handed to the membership + // loop, which re-claims what it holds after this replica has been reaped, + // and to the route that accepts a worker's dial. + Tunnels *cluster.TunnelRegistry shutdownOnce sync.Once } @@ -208,6 +213,26 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } } + // The worker tunnels this replica accepts. It claims as the SAME instance + // ID membership registers under, because that is the ID a peer's Owner + // lookup joins a claim against to decide the owner is alive; two IDs here + // would make every claim this replica writes look like it belongs to a + // replica that does not exist. + tunnels := cluster.NewTunnelRegistry(clusterRegistry, cfg.Distributed.InstanceID) + // Without this the re-claim in the heartbeat loop is dead code: a replica + // stalled long enough to be swept loses the connection rows it owned, and + // nothing would ever write them back, so every other replica would answer + // "not connected" for workers that are connected right here. + // + // Nil when no peer-reachable address could be determined above. There is no + // heartbeat loop to hand it to in that case, and no other replica can reach + // this one anyway; the registry is still built, because it is what the + // tunnel endpoint attaches to and what this replica opens its own streams + // through. + if membership != nil { + membership.SetTunnels(tunnels) + } + // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an // operator can pin placement to a stable name like "production" and have it @@ -499,6 +524,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Cluster: clusterRegistry, Membership: membership, PeerSessions: peerSessions, + Tunnels: tunnels, }, nil } diff --git a/core/http/app.go b/core/http/app.go index a2a3555eafe0..f3e3834bd258 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -593,6 +593,18 @@ func API(application *application.Application) (*echo.Echo, error) { routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, d.PeerSessions.Accept) } + // The worker tunnel, registered unconditionally. Both arguments are nil + // outside distributed mode and the handler refuses every dial then, which + // is what makes registering it always safe; what it buys is the + // route-coverage test walking the route in a plain single-binary + // application, and that test is what holds the rule that an unauthenticated + // dial is refused BEFORE the WebSocket upgrade. + var tunnels *clustersvc.TunnelRegistry + if d := application.Distributed(); d != nil { + tunnels = d.Tunnels + } + routes.RegisterWorkerTunnelRoute(e, registry, tunnels) + // Distributed SSE routes (job progress + agent events via NATS) if d := application.Distributed(); d != nil { if d.Dispatcher != nil { diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go new file mode 100644 index 000000000000..85c48b05277c --- /dev/null +++ b/core/http/endpoints/cluster/connect.go @@ -0,0 +1,191 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "errors" + "net/http" + "strings" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +// ConnectHandler serves the door a worker knocks on: it authenticates the dial +// against the node's OWN stored token, upgrades it to a WebSocket, wraps that as +// a yamux server session and attaches it to the tunnel registry. +// +// The worker dials out and never listens, which is the whole point of the +// tunnel: a worker behind NAT, in another cluster, or on a laptop needs no +// inbound port. It is therefore the yamux CLIENT and this side the SERVER, so +// this side owns the even stream IDs and is the side that opens streams. +// Nothing here accepts streams: in this design the frontend asks and the worker +// answers, so a worker that opened a stream into this session would park on the +// accept backlog rather than be served. +// +// The route is registered in every deployment, including single-binary ones, so +// that the route-coverage test under build tag `auth` walks it. A nil registry +// or a nil tunnel registry therefore has to be a real answer rather than a +// panic; see the 503 below. +// +// It is deliberately absent from auth.RouteFeatureRegistry. That registry gates +// a route on the FEATURES OF AN AUTHENTICATED USER, and there is no user here: +// the dialer is a worker process holding a machine credential, and the global +// auth middleware never runs on this path because it sits under +// auth.ClusterPathPrefix. +func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegistry) echo.HandlerFunc { + // gorilla's default CheckOrigin restricts a browser to same-origin and lets + // a header-less client (which every worker is) through, so the zero value + // is what this link wants. The same choice PeerHandler makes. + upgrader := websocket.Upgrader{} + + return func(c echo.Context) error { + // Everything below happens BEFORE the upgrade, and the order inside it + // is load-bearing. The credential is read first because a dial with no + // Authorization header at all is the anonymous case, and the + // route-coverage test issues exactly that, with no query string: it + // must see 401 rather than a 400 about a missing node id. + token, ok := bearerToken(c.Request()) + if !ok { + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // Not 401. A frontend with no cluster cannot authenticate anybody, and + // answering "unauthorized" would send the operator hunting a token + // problem that does not exist. It is checked after the header so that + // an anonymous dial still gets the 401 the coverage test requires. + 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): + // A node this frontend has never seen. Reported as 401 rather than + // 404 so a caller cannot enumerate node IDs by status code. + 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. This + // is the phase's standing rule in its HTTP form: telling a worker + // its credentials are wrong when the database merely could not be + // read sends it re-registering instead of retrying, and a worker + // that re-registers has thrown away the identity its tunnel and its + // loaded models are 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") + } + + if !authorizedWorker(token, node.TokenHash) { + xlog.Debug("worker tunnel dial presented the wrong token", "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has already written its own failure to the client. + xlog.Debug("worker tunnel upgrade failed", "node", nodeID, "error", err) + return nil + } + + sess, err := yamux.Server(clustersvc.WebsocketConn(ws), nil, nil) + if err != nil { + xlog.Error("Worker tunnel session setup failed", "node", nodeID, "error", err) + _ = ws.Close() + return nil + } + + // From here the connection is hijacked, so no status can reach the + // worker any more: a failure is a closed socket, which is what its + // reconnect loop reads. + epoch, err := tunnels.Attach(c.Request().Context(), nodeID, sess) + if err != nil { + xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "error", err) + _ = sess.Close() + return nil + } + + xlog.Info("Worker tunnel established", "node", nodeID, "remote", ws.RemoteAddr().String()) + // The session outlives this handler, so something other than the + // request goroutine has to notice it die. yamux closes shutdownCh from + // its receive loop the moment the underlying conn fails + // (go-yamux/v5@v5.1.0/session.go:691-695 calling close at + // session.go:297-311), and its default config keepalives every 30s + // (mux.go:73-74), so a worker that vanishes without a FIN is noticed + // too rather than held forever. + go func() { + <-sess.CloseChan() + // The token Attach returned, never a fresh or zero one. Detach + // matches it by EQUALITY: it identifies THIS attachment, so a + // worker that has already re-dialled onto this replica is not + // evicted by its predecessor's teardown. + tunnels.Detach(nodeID, epoch) + xlog.Debug("worker tunnel closed", "node", nodeID) + }() + return nil + } +} + +// bearerToken returns the token from an Authorization: Bearer header, and +// whether one was present at all. +// +// The presence of a credential and its correctness are separate answers on +// purpose: "no credential" is what decides the pre-upgrade 401, and it has to be +// decidable before anything about the node is known. +func bearerToken(r *http.Request) (string, bool) { + // RFC 7235 makes the scheme case-insensitive; the token after it is not. + 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 presented token against the hash stored on the +// node's own row, in constant time. +// +// Against the NODE's stored hash, not the configured registration token, which +// is a deliberate strengthening over how /api/node/ authenticates. A tunnel is a +// durable, multiplexed pipe into a worker; a credential that authorizes every +// worker at once would mean one leak lets an attacker impersonate any worker +// whose ID it can read, and take over that worker's traffic by claiming its +// tunnel. +// +// What that buys TODAY is the mechanism, not yet the isolation, and the +// difference should not be overstated: a worker registers by presenting the +// deployment's registration token, so the hash on its row is that token's hash +// and a leaked registration token still opens a tunnel for a node whose ID the +// attacker knows. What this does rule out is the shortcut of comparing against +// the configured token itself, which is what would have to be unpicked later; +// the day a worker is issued its own secret at registration, this check starts +// isolating workers from each other with no change here. +// +// An empty stored hash authorizes nobody. Such a row exists whenever a worker +// registered against a frontend with no registration token configured, and the +// alternative, letting any token through for those, would make the check depend +// on a setting the dialer can neither see nor be blamed for. +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..7236199171f0 --- /dev/null +++ b/core/http/endpoints/cluster/connect_test.go @@ -0,0 +1,285 @@ +package cluster_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "net/http" + "net/http/httptest" + "net/url" + "strings" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/http/routes" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// workerToken is the secret one worker holds. It is stored on that worker's row +// as a hash, never in a config value, which is the whole point of the check the +// specs below pin: a second worker's token, and the deployment-wide +// registration token, are both wrong for this node. +const workerToken = "worker-1-secret" + +func tokenHash(token string) string { + sum := sha256.Sum256([]byte(token)) + return hex.EncodeToString(sum[:]) +} + +// bearer builds the header a worker dials with. +func bearer(token string) http.Header { + h := http.Header{} + h.Set("Authorization", "Bearer "+token) + return h +} + +// wsConnectURL is the worker tunnel route on a test server, named as nodeID. +func wsConnectURL(s *httptest.Server, nodeID string) string { + return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.ConnectPath + + "?id=" + url.QueryEscape(nodeID) +} + +var _ = Describe("Worker tunnel handler", func() { + var ( + srv *httptest.Server + db *gorm.DB + reg *clustersvc.Registry + tun *clustersvc.TunnelRegistry + nodeID string + ctx context.Context + ) + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + + node := &nodes.BackendNode{ + Name: "worker-1", + Address: "10.0.0.9:50051", + TokenHash: tokenHash(workerToken), + } + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + nodeID = node.ID + Expect(nodeID).ToNot(BeEmpty()) + + reg = clustersvc.NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + tun = clustersvc.NewTunnelRegistry(reg, "me") + + e := echo.New() + routes.RegisterWorkerTunnelRoute(e, nodeReg, tun) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an anonymous dial before upgrading", func() { + // A plain GET, not a WebSocket dial: this is exactly what the + // route-coverage test issues, and a handler that upgrades first answers + // it with gorilla's own 400 handshake failure instead of a 401. + req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+clustersvc.ConnectPath+"?id="+nodeID, nil) + Expect(err).ToNot(HaveOccurred()) + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(resp.Header.Get("Upgrade")).To(BeEmpty(), + "the handler upgraded an unauthenticated dial") + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses a dial that carries no credentials at all", func() { + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), nil) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses a token that is not the node's own", func() { + // The registration token is the credential every worker in the + // deployment holds. Accepting it here would mean one leaked shared + // secret impersonates any worker whose ID an attacker can read. + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer("deployment-registration-token")) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses a dial that names a node it has never seen", func() { + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, "no-such-node"), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses an authenticated dial that names no node", func() { + _, resp, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(srv.URL, "http")+clustersvc.ConnectPath, bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("attaches an authenticated worker and carries bytes to it", func() { + conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + // The worker is the side that dials, so its half of the mux is the + // yamux CLIENT and the frontend's is the server. + workerSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = workerSess.Close() }) + + Eventually(tun.Held, "10s").Should(ConsistOf(nodeID)) + + owner, _, err := reg.OwnerRow(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me"), + "the tunnel was stored without the claim that tells other replicas where it is") + + go func() { + defer GinkgoRecover() + stream, aerr := workerSess.AcceptStream() + if aerr != nil { + return + } + defer func() { _ = stream.Close() }() + buf := make([]byte, 4) + if _, rerr := stream.Read(buf); rerr != nil { + return + } + _, _ = stream.Write(buf) + }() + + stream, err := tun.Open(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + _, err = stream.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + echoed := make([]byte, 4) + _, err = stream.Read(echoed) + Expect(err).ToNot(HaveOccurred()) + Expect(string(echoed)).To(Equal("ping")) + }) + + It("detaches the tunnel and drops its claim when the worker goes away", func() { + conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).ToNot(HaveOccurred()) + Eventually(tun.Held, "10s").Should(ConsistOf(nodeID)) + + Expect(conn.Close()).To(Succeed()) + + Eventually(tun.Held, "10s").Should(BeEmpty(), + "a dead tunnel is still held here, so every dialer routed to this replica gets a socket that carries nothing") + Eventually(func() error { + _, _, err := reg.OwnerRow(ctx, nodeID) + return err + }, "10s").Should(MatchError(clustersvc.ErrNoConnection), + "the claim outlived the socket, so this replica keeps being named the owner of a worker it no longer holds") + }) + + It("reports a lookup failure as a failure, not as a refusal", func() { + // ErrNotOwner, 401 and 404 are all ANSWERS. A database that cannot be + // read is none of them: telling a worker its credentials are wrong when + // the frontend simply could not look them up sends it re-registering + // instead of retrying. + Expect(db.Exec(`DROP TABLE backend_nodes CASCADE`).Error).To(Succeed()) + + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError)) + }) +}) + +var _ = Describe("Worker tunnel handler without distributed mode", func() { + // The route is registered in every deployment so that the route-coverage + // test sees it, which is what pins the reject-before-upgrade rule. With no + // node registry there is nothing to authenticate against, so it must refuse + // every dial rather than publish an unauthenticated multiplexer. + var srv *httptest.Server + + BeforeEach(func() { + e := echo.New() + routes.RegisterWorkerTunnelRoute(e, nil, nil) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an anonymous dial with 401", func() { + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.ConnectPath, nil) + Expect(err).ToNot(HaveOccurred()) + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("tells a credentialed worker the frontend has no cluster, rather than rejecting it", func() { + // A single-binary frontend cannot authenticate anybody, and saying + // "unauthorized" would send an operator hunting a token problem that + // does not exist. + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, "w1"), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + }) +}) + +var _ = Describe("Worker tunnel auth coverage", func() { + // The same argument the peer link's coverage specs make: the tunnel route + // authenticates a worker against its own stored token, not a session, so it + // only works while its path sits under the prefix the global auth + // middleware exempts. + var srv *httptest.Server + + BeforeEach(func() { + e := echo.New() + // A nil DB with one legacy API key is the cheapest configuration that + // turns the middleware ON without a database. + e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}})) + routes.RegisterWorkerTunnelRoute(e, nil, nil) + e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an uncredentialed request to a route outside the cluster prefix", func() { + resp, err := http.Get(srv.URL + "/api/nodes") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), + "the global auth middleware is not actually guarding this server, so the assertion below would prove nothing") + }) + + It("lets a worker dial reach the handler, which is the only thing that can authenticate it", func() { + // The worker's token is not one of the API keys the middleware knows, + // so a 503 from the handler's own no-cluster check can only mean the + // request was let through by the middleware. + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.ConnectPath+"?id=w1", nil) + Expect(err).ToNot(HaveOccurred()) + req.Header.Set("Authorization", "Bearer "+workerToken) + + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable), + "a worker dial must reach the handler; 401 here means the tunnel route left the auth-exempt prefix %q", auth.ClusterPathPrefix) + }) +}) diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go index 3e462db1cd75..458122d036e1 100644 --- a/core/http/routes/cluster.go +++ b/core/http/routes/cluster.go @@ -3,6 +3,7 @@ package routes import ( clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" "github.com/labstack/echo/v4" "github.com/libp2p/go-yamux/v5" @@ -23,3 +24,22 @@ import ( func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { e.GET(clustersvc.PeerPath, clusterep.PeerHandler(token, onPeer)) } + +// RegisterWorkerTunnelRoute registers the endpoint a worker dials to open its +// tunnel. registry authenticates the dial against the node's own stored token; +// tunnels is what the resulting session is attached to. +// +// Unlike the peer link this is registered in EVERY deployment, single-binary +// ones included, and both arguments may be nil there. Two reasons. The handler +// fails closed without a registry, since a token can only be checked against a +// node row and there are none; and being registered unconditionally is what +// puts the route in front of the route-coverage test under build tag `auth`, +// which is the thing that holds the reject-before-upgrade rule in place. A +// route registered only in distributed mode is invisible to that test. +// +// Like the peer link, it carries no auth middleware and derives its path from +// core/services/cluster's own constant, so the handler and the worker's dialler +// cannot end up on different paths. +func RegisterWorkerTunnelRoute(e *echo.Echo, registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegistry) { + e.GET(clustersvc.ConnectPath, clusterep.ConnectHandler(registry, tunnels)) +} diff --git a/core/services/cluster/tunnel.go b/core/services/cluster/tunnel.go index 36f7c91d3d4b..66baa475fdce 100644 --- a/core/services/cluster/tunnel.go +++ b/core/services/cluster/tunnel.go @@ -15,6 +15,17 @@ import ( "github.com/mudler/xlog" ) +// ConnectPath is the route a worker dials to open its tunnel, and the route the +// HTTP layer registers the handler on. It lives here, beside the registry that +// holds what the dial produces, for the reason PeerPath does: the HTTP +// endpoints package imports this one, never the other way round. +// +// The literal is spelled out rather than derived from auth.ClusterPathPrefix +// because importing core/http/auth is exactly the dependency this package must +// not have. A spec in the endpoints package, which can see both, holds the two +// from drifting apart. +const ConnectPath = "/api/cluster/connect" + // ErrNotOwner reports that this replica does not hold the tunnel for a node. // // It is a ROUTING fact and nothing else: some other replica may hold that diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 3a6f9768f669..788f8dee2d55 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -100,6 +100,22 @@ environment: The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_REGISTRATION_TOKEN`, the same shared secret workers register with. Replicas that disagree about it cannot link. A replica that stops heartbeating for 30 seconds is dropped from the table by the others, along with the worker-connection rows it owned. +### Worker tunnels + +A worker can open one long-lived, multiplexed tunnel to the frontend instead of listening on a port of its own. It dials `GET /api/cluster/connect?id=`, the connection is upgraded to a WebSocket, and every subsequent request the frontend makes to that worker travels as a stream inside it. Nothing dials *into* the worker, so a worker behind NAT, in another Kubernetes cluster or on a laptop needs no inbound port and no reachable address. + +The dial is authenticated against **that node's own stored token**, not the deployment-wide registration token: the frontend hashes the presented bearer token and compares it with the hash recorded on the node's row at registration. A worker that presents a token belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity. + +One honest caveat about that today: a worker currently registers by presenting the deployment's registration token, so the hash stored on its row *is* that token's hash, and a leaked registration token plus a known node ID still gets a tunnel. What the tunnel endpoint does not do is trust the configured token directly, so when workers are issued their own secrets at registration the isolation becomes real with no change to this route. + +The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims every tunnel it still holds as soon as it re-registers. + +| Method | Path | Description | +|--------|------|-------------| +| `GET` | `/api/cluster/connect?id=` | Worker opens its multiplexed tunnel (`Authorization: Bearer `) | + +The route is exempt from the normal session/API-key authentication (it authenticates itself, like `/api/cluster/peer`) and is registered in every deployment. Outside distributed mode there is no node table to check a token against, so it answers `503`. + ### The model load deadline scales with the checkpoint The `LoadModel` deadline starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint read and pipeline init. That work is proportional to the bytes on disk, which makes any fixed deadline a model-size cliff rather than a timeout: a 70 GB video checkpoint on a Jetson Thor worker failed reproducibly against the old fixed 5m default (`rpc error: code = DeadlineExceeded` after 953.5s of wall clock, roughly 11m of which was backend install and staging), and simply raising the constant would only move the cliff to the next larger model while making a genuinely wedged *small* model hang for the whole inflated duration. @@ -486,6 +502,8 @@ Used by workers themselves (registration, heartbeat, etc.). Authenticated via th | `GET` | `/api/node/:id/models` | Query own loaded models | | `DELETE` | `/api/node/:id` | Deregister self | +The worker tunnel at `GET /api/cluster/connect` is also worker-facing but is authenticated differently: against the node's own stored token rather than the shared registration token. See [Worker tunnels](#worker-tunnels). + ### `/api/nodes/` - Admin management Used by the WebUI and admin API consumers. Requires admin authentication. From 48ece89c638fbe2da112e3d3822aa461722d9eaf Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 10:36:10 +0000 Subject: [PATCH 23/42] fix(cluster): harden the worker tunnel, and stop starting a database per spec Review follow-up. Twelve findings, none blocking, grouped here by what they protect. Panics. The handler now recovers between the WebSocket upgrade and the hand-off, the way the peer link next door already did: net/http recovers the panic but leaves the hijacked socket open, so without this a worker keeps a session this replica has no entry for and will never detach. The claim gate in Attach and reclaimOne is now released with defer, so a panic under Claim cannot wedge one node's gate for the life of the process. SetTunnels gained the nil-receiver guard its sibling Stop has. Operability. A deployment with no registration token stores an empty token_hash on every worker, so every tunnel dial 401s forever on a frontend that looks correctly configured. That now warns at startup, logs its own line rather than sharing the "wrong token" one, and is stated in the docs together with the fact that setting the token later needs the workers to register again. Authorization. A node still awaiting admin approval is refused with 403. The rest of /api/node/ gates on nothing, but the two places that hand a node something durable, its API key and its NATS credential, both refuse a pending one, and a tunnel is that kind of grant. Draining and unhealthy nodes keep their tunnels on purpose. Comments that claimed more than the code. The global auth middleware does run on this path and then declines to reject; the future per-node secret only lands without a change here if it lands in TokenHash; the empty-hash guard is defensive rather than deciding; ClusterPathPrefix is no longer only replica-to-replica; the docs no longer say a reaped replica re-claims unconditionally. And the test harness. SetupTestDB started a PostgreSQL container per BeforeEach with a readiness deadline it asserted on, which is one chance per spec to fail one spec inside its setup, anywhere, never twice in the same place: the shape of the flake seen twice here and never reproduced. It now starts one container per process and creates a database per call, which is the pattern tests/e2e already proved. Isolation is unchanged and is now asserted for the first time. All 69 call sites are untouched; the eleven consumer packages run 1404 specs green, and jobs went from 34.3s to 3.3s, agents from 13.8s to 1.9s, cluster from 97.4s to 37.5s. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] --- core/http/app.go | 13 ++ core/http/auth/public_routes.go | 11 +- core/http/endpoints/cluster/connect.go | 87 +++++++++- core/http/endpoints/cluster/connect_test.go | 89 ++++++++++ core/services/cluster/membership.go | 8 + core/services/cluster/tunnel.go | 99 +++++++---- core/services/cluster/tunnel_test.go | 75 ++++++++ core/services/testutil/testdb.go | 183 ++++++++++++++++++-- core/services/testutil/testdb_test.go | 52 ++++++ docs/content/features/distributed-mode.md | 10 +- 10 files changed, 559 insertions(+), 68 deletions(-) create mode 100644 core/services/testutil/testdb_test.go diff --git a/core/http/app.go b/core/http/app.go index f3e3834bd258..8f4f34912d24 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -602,6 +602,19 @@ func API(application *application.Application) (*echo.Echo, error) { var tunnels *clustersvc.TunnelRegistry if d := application.Distributed(); d != nil { tunnels = d.Tunnels + if distCfg.RegistrationToken == "" { + // A separate warning from the peer link's, for the same missing + // knob, because the broken thing is different and an operator has + // to be told both. A worker registering against a frontend with no + // registration token sends no token, so RegisterNodeEndpoint stores + // an empty token_hash, and the tunnel authenticates against exactly + // that column: every dial 401s, forever, on a deployment that looks + // correctly configured. Fixing it needs the token set AND the + // workers re-registered, which is why it is worth saying at boot + // rather than leaving in the 401s. + xlog.Warn("Worker tunnels will refuse every dial: no registration token is configured, so no worker has a stored token to authenticate against", + "route", clustersvc.ConnectPath, "knob", "LOCALAI_REGISTRATION_TOKEN") + } } routes.RegisterWorkerTunnelRoute(e, registry, tunnels) diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index 4a8fcff01cac..2c1e1dc3cdd1 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -74,9 +74,14 @@ func isPublicRoute(method, path string) bool { return false } -// ClusterPathPrefix is the replica-to-replica namespace. Its handlers check the -// cluster token in the Authorization header themselves, so the check below lets -// them through the global session middleware. +// ClusterPathPrefix is the machine-to-machine cluster namespace. It carries two +// different trust relationships, on two different credentials: the +// replica-to-replica peer link, which checks the shared cluster token, and the +// worker-to-frontend tunnel, which checks the dialing node's own stored token +// hash. What they have in common is the only thing this prefix asserts, that +// each handler checks its own Authorization header, so the check below lets them +// through the global session middleware rather than rejecting a caller that has +// no session and no user. // // The cluster routes do NOT derive their paths from this constant: they are // registered from core/services/cluster's own literal, because that package diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go index 85c48b05277c..e76450692fd9 100644 --- a/core/http/endpoints/cluster/connect.go +++ b/core/http/endpoints/cluster/connect.go @@ -37,10 +37,16 @@ import ( // panic; see the 503 below. // // It is deliberately absent from auth.RouteFeatureRegistry. That registry gates -// a route on the FEATURES OF AN AUTHENTICATED USER, and there is no user here: -// the dialer is a worker process holding a machine credential, and the global -// auth middleware never runs on this path because it sits under -// auth.ClusterPathPrefix. +// a route on the FEATURES OF AN AUTHENTICATED USER, resolved from auth.GetUser, +// and there is no user here: the dialer is a worker process holding a machine +// credential. +// +// The global auth middleware does RUN on this path; what it does not do is +// reject. It attempts session, bearer and legacy-key authentication first, so it +// may even have set auth_user from a worker token that happens to match an API +// key, and then core/http/auth/middleware.go:90 lets the request through +// because usesAlternativeAuthentication reports the path as one whose +// credentials its own route checks. Nothing here reads what it set. func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegistry) echo.HandlerFunc { // gorilla's default CheckOrigin restricts a browser to same-origin and lets // a header-less client (which every worker is) through, so the zero value @@ -62,6 +68,14 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi // answering "unauthorized" would send the operator hunting a token // problem that does not exist. It is checked after the header so that // an anonymous dial still gets the 401 the coverage test requires. + // + // Only the registry half is covered by a spec. The two are read from one + // application.Distributed() in core/http/app.go and initDistributed + // returns an error rather than a partial struct, so a non-nil registry + // beside a nil tunnel registry is unreachable and no spec constructs it; + // the second half is defence against a future wiring that splits them, + // where the cost would be a nil dereference in Attach after the + // connection is already hijacked. if registry == nil || tunnels == nil { return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled") } @@ -89,11 +103,42 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi return echo.NewHTTPError(http.StatusInternalServerError, "node lookup failed") } + // Split from the mismatch below because they are different operator + // problems with different fixes. An empty stored hash means the node + // registered against a frontend with no LOCALAI_REGISTRATION_TOKEN, so + // no worker on that deployment can ever tunnel until one is set and the + // workers re-register; a mismatch means one worker holds the wrong + // secret. One log line for both leaves an operator reading "wrong token" + // while every worker fails identically. + if node.TokenHash == "" { + xlog.Warn("Refusing a worker tunnel: this node has no stored token, which means it registered with no registration token configured", + "node", nodeID, "knob", "LOCALAI_REGISTRATION_TOKEN") + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } if !authorizedWorker(token, node.TokenHash) { xlog.Debug("worker tunnel dial presented the wrong token", "node", nodeID) return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") } + // Authenticated but not authorised, so 403 rather than 401: the fix is an + // admin approving the node, not a different credential, and answering + // 401 would send an operator looking at tokens. + // + // Only StatusPending is refused. The rest of /api/node/ self-service + // gates on nothing at all, but the two places that hand a node something + // DURABLE both refuse a pending one: the agent worker's API key + // (core/http/endpoints/localai/nodes.go:224) and its NATS credential + // (nodes.go:293). A tunnel is that kind of grant, not a heartbeat: it is + // a standing pipe into the worker recorded in node_connections and + // relayed to by every other replica. Draining and unhealthy nodes keep + // their tunnels on purpose; draining means finish what you have, and a + // node marked unhealthy for missed heartbeats needs the pipe to recover + // through. + 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") + } + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) if err != nil { // Upgrade has already written its own failure to the client. @@ -108,6 +153,23 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi return nil } + // The same guard PeerHandler carries, for the same reason and one more. + // net/http recovers a panic from this goroutine but does not close the + // hijacked connection, and middleware.Recover does not either, so a + // panic below would leave the worker holding a live session this replica + // has no entry for and will never detach. The extra reason here is that + // Attach does database work: a panic inside it, with the session left + // open, is a tunnel nothing can reach and nothing will clean up. + // + // It re-panics rather than swallowing. Whatever it caught is a bug, and + // the recovery middleware above is what should report it. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() + // From here the connection is hijacked, so no status can reach the // worker any more: a failure is a closed socket, which is what its // reconnect loop reads. @@ -176,12 +238,19 @@ func bearerToken(r *http.Request) (string, bool) { // attacker knows. What this does rule out is the shortcut of comparing against // the configured token itself, which is what would have to be unpicked later; // the day a worker is issued its own secret at registration, this check starts -// isolating workers from each other with no change here. +// isolating workers from each other with no change here, PROVIDED the secret +// lands in TokenHash. The one per-node secret LocalAI mints today, the agent +// worker's api_token, does not: provisionAgentWorkerKey writes an auth.User and +// an auth.APIKey referenced by node.AuthUserID / node.APIKeyID and never touches +// this column. If the next task follows that precedent instead, this comparison +// is what has to change. // -// An empty stored hash authorizes nobody. Such a row exists whenever a worker -// registered against a frontend with no registration token configured, and the -// alternative, letting any token through for those, would make the check depend -// on a setting the dialer can neither see nor be blamed for. +// The empty-hash guard is defensive rather than deciding: a stored hash is +// hex-encoded SHA-256, so 64 bytes or nothing, and ConstantTimeCompare already +// returns 0 on a length mismatch (crypto/internal/fips140/subtle/constant_time.go:17-20 +// returns 0 outright when the lengths differ). It is kept because a reader should not have to +// derive "an unregistered node authorizes nobody" from a length rule, and +// because the caller logs that case separately. 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 7236199171f0..072f69d0226c 100644 --- a/core/http/endpoints/cluster/connect_test.go +++ b/core/http/endpoints/cluster/connect_test.go @@ -194,6 +194,45 @@ var _ = Describe("Worker tunnel handler", func() { "the claim outlived the socket, so this replica keeps being named the owner of a worker it no longer holds") }) + It("refuses a node whose row carries no stored token", func() { + // A deployment with no registration token configured produces exactly + // these rows: the worker sends no token, so registration stores none, + // and there is nothing here to authenticate against. + Expect(db.Exec(`UPDATE backend_nodes SET token_hash = '' WHERE id = ?`, nodeID).Error).To(Succeed()) + + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("refuses a node that is still awaiting admin approval", func() { + // Approval is what gates a node's participation, and a tunnel is a + // standing pipe recorded in node_connections, not a heartbeat. 403 and + // not 401: the credential is right, the authorisation is missing, and + // the fix is an admin rather than a different token. + Expect(db.Exec(`UPDATE backend_nodes SET status = ? WHERE id = ?`, nodes.StatusPending, nodeID).Error).To(Succeed()) + + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).To(HaveOccurred()) + Expect(resp).ToNot(BeNil()) + Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) + Expect(tun.Held()).To(BeEmpty()) + }) + + It("still admits a draining node, which has work to finish", func() { + // Only pending is refused. Draining means "start nothing new", not + // "lose the pipe your in-flight requests travel on", and a node marked + // unhealthy for missed heartbeats needs the tunnel to recover through. + Expect(db.Exec(`UPDATE backend_nodes SET status = ? WHERE id = ?`, nodes.StatusDraining, nodeID).Error).To(Succeed()) + + conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Eventually(tun.Held, "10s").Should(ConsistOf(nodeID)) + }) + It("reports a lookup failure as a failure, not as a refusal", func() { // ErrNotOwner, 401 and 404 are all ANSWERS. A database that cannot be // read is none of them: telling a worker its credentials are wrong when @@ -208,6 +247,56 @@ var _ = Describe("Worker tunnel handler", func() { }) }) +var _ = Describe("Worker tunnel handler when the attach panics", func() { + // net/http recovers a panic from the request goroutine but does NOT close a + // hijacked connection, so without the handler's own recover the worker keeps + // a live session this replica has no entry for and will never detach: its + // opens fill yamux's 256-deep backlog and then hang with no error. The panic + // is injected through a real path rather than a fake one, a registry built + // over no database at all, which is what Attach's first database call + // dereferences. + var ( + srv *httptest.Server + db *gorm.DB + nodeID string + ) + + BeforeEach(func() { + ctx := context.Background() + db = testutil.SetupTestDB() + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + node := &nodes.BackendNode{Name: "worker-1", Address: "10.0.0.9:50051", TokenHash: tokenHash(workerToken)} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + nodeID = node.ID + + e := echo.New() + routes.RegisterWorkerTunnelRoute(e, nodeReg, clustersvc.NewTunnelRegistry(nil, "me")) + srv = httptest.NewServer(e) + DeferCleanup(func() { + // A hijacked connection the handler never closed would park Close + // forever, turning the assertion below into a suite hang. + srv.CloseClientConnections() + srv.Close() + }) + }) + + It("closes the worker's session instead of stranding it", func() { + conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + workerSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = workerSess.Close() }) + + // Asserting on OpenStream would hang rather than fail: yamux only + // acknowledges a stream once the peer accepts it, and the leak this + // pins is precisely that nobody ever will. + Eventually(workerSess.IsClosed, "10s").Should(BeTrue()) + }) +}) + var _ = Describe("Worker tunnel handler without distributed mode", func() { // The route is registered in every deployment so that the route-coverage // test sees it, which is what pins the reject-before-upgrade rule. With no diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go index 4a5f3e19ccab..6a8650927b15 100644 --- a/core/services/cluster/membership.go +++ b/core/services/cluster/membership.go @@ -81,7 +81,15 @@ func NewMembership(reg *Registry, id, addr, version string) *Membership { // It is a setter rather than a constructor argument because the tunnel registry // is what the tunnel endpoint is built on, and that is wired after membership // is already running. +// +// Safe on a nil receiver, like Stop. This package deliberately produces a nil +// *Membership (core/application/distributed.go leaves it nil when no +// peer-reachable address can be derived), so a setter that panicked on one +// would be a trap for the next caller rather than an impossibility. func (m *Membership) SetTunnels(t *TunnelRegistry) { + if m == nil { + return + } m.mu.Lock() defer m.mu.Unlock() m.tunnels = t diff --git a/core/services/cluster/tunnel.go b/core/services/cluster/tunnel.go index 66baa475fdce..65adc4aad4ca 100644 --- a/core/services/cluster/tunnel.go +++ b/core/services/cluster/tunnel.go @@ -182,18 +182,33 @@ func (t *TunnelRegistry) Attach(ctx context.Context, nodeID string, sess *yamux. return 0, fmt.Errorf("attaching tunnel for node %q: %w", nodeID, err) } - epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) + // The gated part is a closure so its release can be DEFERRED while the + // session close below still happens outside the gate. Releasing on each + // return path instead leaves one way out uncovered: a panic. Claim does + // database work, and a panic anywhere under it would leave this node's gate + // closed for the life of the process, so every later Attach or Reclaim for + // that worker would block in enterClaim until its own context expired. The + // caller's recover would report the panic and the worker would look + // permanently unable to reconnect, with nothing linking the two. + 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{sess: sess, token: epoch, claim: epoch} + t.mu.Unlock() + return epoch, nil + }() if err != nil { - t.leaveClaim(nodeID) return 0, err } - t.mu.Lock() - previous := t.tunnels[nodeID] - t.tunnels[nodeID] = &heldTunnel{sess: sess, token: epoch, claim: epoch} - t.mu.Unlock() - t.leaveClaim(nodeID) - // Closed after the gate is released, not under it. The gate is justified by // being held for one claim round trip, and closing a session is not that: // yamux closes the underlying conn and then waits for both its send and @@ -381,38 +396,48 @@ func (t *TunnelRegistry) reclaimOne(ctx context.Context, nodeID string) error { return fmt.Errorf("re-claiming node %q: %w", nodeID, err) } - t.mu.Lock() - tunnel, ok := t.tunnels[nodeID] - t.mu.Unlock() - if !ok { - // Detached between Reclaim listing the nodes and this gate. Nothing was - // claimed, so there is nothing to undo. - t.leaveClaim(nodeID) - return errTunnelNotReclaimed - } - if tunnel.sess.IsClosed() { - xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID) - t.leaveClaim(nodeID) - return errTunnelNotReclaimed - } + // Gated part in a closure so the release is DEFERRED, for the reason Attach + // gives: a panic under Claim would otherwise wedge this node's gate for the + // life of the process. The trailing Release still runs outside the gate. + var epoch int64 + var installed bool + if err := func() error { + defer t.leaveClaim(nodeID) - epoch, err := t.reg.Claim(ctx, nodeID, t.selfID) - if err != nil { - t.leaveClaim(nodeID) - return err - } + t.mu.Lock() + tunnel, ok := t.tunnels[nodeID] + t.mu.Unlock() + if !ok { + // Detached between Reclaim listing the nodes and this gate. Nothing + // was claimed, so there is nothing to undo. + return errTunnelNotReclaimed + } + if tunnel.sess.IsClosed() { + xlog.Debug("skipping re-claim of a worker tunnel whose session is closed", "node", nodeID) + return errTunnelNotReclaimed + } - t.mu.Lock() - current, installed := t.tunnels[nodeID] - if installed { - // current is necessarily the entry read above: the gate is still held, - // and Attach and reclaimOne are the only writers that install one. The - // identity is therefore not re-checked; the case that IS reachable is - // the entry being gone, because Detach is not gated. - current.claim = epoch + 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 necessarily the entry read above: the gate is still + // held, and Attach and reclaimOne are the only writers that install + // one. The identity is therefore not re-checked; the case that IS + // reachable is the entry being gone, because Detach is not gated. + current.claim = epoch + } + t.mu.Unlock() + return nil + }(); err != nil { + return err } - t.mu.Unlock() - t.leaveClaim(nodeID) if installed { return nil diff --git a/core/services/cluster/tunnel_test.go b/core/services/cluster/tunnel_test.go index 67993130ee54..acc81d81813f 100644 --- a/core/services/cluster/tunnel_test.go +++ b/core/services/cluster/tunnel_test.go @@ -716,3 +716,78 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu Consistently(stored, 2*cluster.InstanceHeartbeat, time.Second).Should(Equal(epoch)) }) }) + +var _ = Describe("The worker tunnel registry's claim gate", func() { + // These specs need no database. They pin what happens when the database + // call under the gate does not return normally, which is the one exit a + // release-on-every-return-path cannot cover. + // + // The panic is produced by the production code itself: a registry with no + // *Registry behind it dereferences nothing at the gate and then panics + // inside Claim, which is exactly where a real one does its work. + var tun *cluster.TunnelRegistry + + BeforeEach(func() { + tun = cluster.NewTunnelRegistry(nil, "me") + }) + + // attachPanics runs one Attach that is expected to panic, swallowing the + // panic so the spec can go on to ask what state it left behind. + attachPanics := func(ctx context.Context, nodeID string) { + defer GinkgoRecover() + defer func() { _ = recover() }() + frontend, _ := workerTunnel() + _, _ = tun.Attach(ctx, nodeID, frontend) + Fail("Attach was expected to panic inside Claim, so this spec is no longer testing what it claims") + } + + It("frees the node's gate when the claim panics", func() { + attachPanics(context.Background(), "w1") + + // A wedged gate is indistinguishable from a slow one except by waiting, + // so the second attempt is given a deadline. Reaching Claim means + // panicking again; returning a context error means it never got past + // enterClaim and this worker could never reconnect to this replica. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + reached := make(chan any, 1) + go func() { + defer GinkgoRecover() + defer func() { reached <- recover() }() + frontend, _ := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).To(MatchError(context.DeadlineExceeded), + "the gate for this node was never released, so every later dial from it blocks until its own context expires") + }() + Eventually(reached, "5s").Should(Receive(Not(BeNil())), + "the second Attach did not reach Claim, so the panicking one left the gate closed") + }) + + It("leaves another node's gate alone", func() { + // The gate is per node so that one wedged worker cannot stop the rest; + // this holds that property against the panic path too. + attachPanics(context.Background(), "w1") + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + reached := make(chan any, 1) + go func() { + defer GinkgoRecover() + defer func() { reached <- recover() }() + frontend, _ := workerTunnel() + _, _ = tun.Attach(ctx, "w2", frontend) + }() + Eventually(reached, "5s").Should(Receive(Not(BeNil()))) + }) +}) + +var _ = Describe("Membership.SetTunnels", func() { + It("is safe on a nil receiver, like Stop", func() { + // A nil *Membership is a value this codebase deliberately produces when + // no peer-reachable address can be derived, so the asymmetry with Stop + // would be a trap for the next caller. + var m *cluster.Membership + Expect(func() { m.SetTunnels(cluster.NewTunnelRegistry(nil, "me")) }).ToNot(Panic()) + }) +}) diff --git a/core/services/testutil/testdb.go b/core/services/testutil/testdb.go index 80e511201b7d..755589cd438c 100644 --- a/core/services/testutil/testdb.go +++ b/core/services/testutil/testdb.go @@ -2,7 +2,11 @@ package testutil import ( "context" + "fmt" + "net/url" "runtime" + "sync" + "sync/atomic" "time" "github.com/testcontainers/testcontainers-go" @@ -16,27 +20,176 @@ import ( . "github.com/onsi/gomega" ) -// SetupTestDB creates a fresh PostgreSQL 16 container and returns a gorm.DB. -// The container is cleaned up via DeferCleanup when the test completes. +// One PostgreSQL container per test PROCESS, not per spec, with a database per +// SetupTestDB call. +// +// Starting a container per spec was both slow and flaky. Slow because a +// postgres:16 start is seconds and the packages behind this helper hold several +// hundred specs; flaky because every start was a fresh chance to miss the +// readiness deadline, and a miss lands in the caller's BeforeEach as a failure +// of whichever spec happened to be running. That is the exact shape of the +// intermittent single-spec failure seen twice in this package and never +// reproduced: one spec of many, no pattern, never twice in the same place. +// Starting the container once per process leaves one chance to miss it instead +// of one per spec, and moves that chance onto a deadline that only has to be met +// while nothing else is competing for the machine. +// +// Isolation is unchanged and is what callers actually depend on: each call still +// hands back an empty database that no other spec can see. The database is +// dropped when the spec that asked for it ends. Advisory locks, sequences and +// extensions are all per-database in PostgreSQL, so nothing the packages behind +// this helper rely on leaks between specs. +// +// This mirrors the pattern already proven in tests/e2e/distributed +// (testhelpers_test.go), which is where the argument and the measurements come +// from. +// +// One container per process rather than one shared across `ginkgo -p` workers is +// deliberate: parallel Ginkgo processes are separate OS processes, each gets its +// own container, and nothing has to coordinate database names across them. +var ( + sharedOnce sync.Once + sharedPG *tcpostgres.PostgresContainer + sharedDSN string + sharedErr error + + // dbCounter makes each database name unique within this process. The + // container is not shared across processes, so a process-local counter is + // enough. + dbCounter atomic.Int64 +) + +// The container outlives every spec, so its teardown belongs to the suite. This +// registers one AfterSuite in every suite that imports this package, which is +// every suite that could have started a container; it is a no-op in the ones +// that never call SetupTestDB. +// +// Package-level rather than something callers have to remember: a helper whose +// cleanup depends on 56 test files each declaring a hook is a helper that leaks +// containers the first time someone forgets. Registration happens during package +// initialisation, which is before RunSpecs, so Ginkgo is still building its tree. +var _ = AfterSuite(func() { + if sharedPG == nil { + return + } + // Best-effort: a failed terminate must not fail a suite whose specs all + // passed. Testcontainers' reaper removes it in that case. + _ = sharedPG.Terminate(context.Background()) +}) + +// sharedPostgres returns the DSN of this process's PostgreSQL container, +// starting it on first use. +// +// The error is remembered rather than only asserted inside the sync.Once: an +// assertion there fails the one spec that happened to be first, and every later +// spec would then find a nil container and fail for some unrelated-looking +// reason. Re-asserting the stored error makes every affected spec say the same +// true thing. +func sharedPostgres() string { + GinkgoHelper() + + sharedOnce.Do(func() { + ctx := context.Background() + sharedPG, sharedErr = tcpostgres.Run(ctx, "postgres:16", + tcpostgres.WithDatabase("testdb"), + tcpostgres.WithUsername("test"), + tcpostgres.WithPassword("test"), + // The deadline is per process now, not per spec, so it is generous + // on purpose: it is paid once, and the cost of missing it is a + // whole suite rather than one spec. + testcontainers.WithWaitStrategyAndDeadline(120*time.Second, + wait.ForLog("database system is ready to accept connections").WithOccurrence(2)), + ) + if sharedErr != nil { + return + } + sharedDSN, sharedErr = sharedPG.ConnectionString(ctx, "sslmode=disable") + }) + + Expect(sharedErr).ToNot(HaveOccurred(), "the suite's PostgreSQL container could not be started") + return sharedDSN +} + +// SetupTestDB returns a gorm.DB on a PostgreSQL database created for the calling +// spec. The database is dropped, and its connection pool closed, when the spec +// ends. func SetupTestDB() *gorm.DB { + GinkgoHelper() if runtime.GOOS == "darwin" { Skip("testcontainers requires Docker, not available on macOS CI") } - ctx := context.Background() - pgC, err := tcpostgres.Run(ctx, "postgres:16", - tcpostgres.WithDatabase("testdb"), - tcpostgres.WithUsername("test"), - tcpostgres.WithPassword("test"), - testcontainers.WithWaitStrategyAndDeadline(60*time.Second, - wait.ForLog("database system is ready to accept connections").WithOccurrence(2)), - ) - Expect(err).ToNot(HaveOccurred()) - DeferCleanup(func() { pgC.Terminate(context.Background()) }) - connStr, err := pgC.ConnectionString(ctx, "sslmode=disable") - Expect(err).ToNot(HaveOccurred()) - db, err := gorm.Open(postgres.Open(connStr), &gorm.Config{ + + dsn := sharedPostgres() + name := fmt.Sprintf("testdb_%d", dbCounter.Add(1)) + + // Scoped so a failed CREATE cannot leak the pool: the assertion panics out + // of this function, and a leaked pool per failing spec exhausts the + // server's connection limit for every spec after it. + // + // CREATE and DROP DATABASE cannot run against the target database itself, + // so both go through a short-lived connection to the container's own + // maintenance database. + func() { + admin := openPool(dsn) + defer closePool(admin) + Expect(admin.Exec(fmt.Sprintf("CREATE DATABASE %q", name)).Error).To(Succeed()) + }() + + db, err := gorm.Open(postgres.Open(replaceDBName(dsn, name)), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) Expect(err).ToNot(HaveOccurred()) + + DeferCleanup(func() { + // The caller's own DeferCleanups were registered later and so run + // first, which is what lets a spec keep using this handle in its + // teardown. + closePool(db) + + drop, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) + if err != nil { + // Reported, never asserted. A cleanup that fails the spec turns one + // database hiccup into a failure that buries whatever the spec was + // actually about. + AddReportEntry("drop test database skipped", fmt.Sprintf("%s: %v", name, err)) + return + } + defer closePool(drop) + // FORCE terminates whatever connections the spec left open, including + // any a background goroutine is still holding (PostgreSQL 13+). + if err := drop.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %q WITH (FORCE)", name)).Error; err != nil { + AddReportEntry("drop test database failed", fmt.Sprintf("%s: %v", name, err)) + } + }) + + return db +} + +// openPool connects to dsn with logging off. Used for the short-lived +// maintenance connections only; the database a spec is handed keeps gorm's +// silent logger so a caller can still swap it. +func openPool(dsn string) *gorm.DB { + GinkgoHelper() + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) + Expect(err).ToNot(HaveOccurred()) return db } + +func closePool(db *gorm.DB) { + if db == nil { + return + } + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } +} + +// replaceDBName swaps the database component of a DSN, preserving credentials, +// host, port and query parameters. +func replaceDBName(dsn, name string) string { + GinkgoHelper() + u, err := url.Parse(dsn) + Expect(err).ToNot(HaveOccurred()) + u.Path = "/" + name + return u.String() +} diff --git a/core/services/testutil/testdb_test.go b/core/services/testutil/testdb_test.go new file mode 100644 index 000000000000..d9dd173eb20f --- /dev/null +++ b/core/services/testutil/testdb_test.go @@ -0,0 +1,52 @@ +package testutil_test + +import ( + "testing" + + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestTestutil(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Test Utilities Suite") +} + +// The container is shared per process now, so the isolation callers depend on +// comes from a database per call rather than from a server per call. That is +// the property 69 call sites across eleven packages assume without saying so, +// and nothing else in the tree asserts it. +var _ = Describe("SetupTestDB", func() { + type row struct { + ID int + } + + It("hands back a database no other caller can see into", func() { + first := testutil.SetupTestDB() + second := testutil.SetupTestDB() + + Expect(first.Exec(`CREATE TABLE isolation_probe (id int)`).Error).To(Succeed()) + Expect(first.Exec(`INSERT INTO isolation_probe VALUES (1)`).Error).To(Succeed()) + + var found []row + err := second.Raw(`SELECT id FROM isolation_probe`).Scan(&found).Error + Expect(err).To(HaveOccurred(), + "two SetupTestDB calls landed on the same database, so every spec can now see every other spec's rows") + + // The second database must also be usable, not merely different: an + // isolation check that passed because the second handle was broken + // would prove nothing. + Expect(second.Exec(`CREATE TABLE isolation_probe (id int)`).Error).To(Succeed()) + }) + + It("hands back an empty database", func() { + db := testutil.SetupTestDB() + var tables int64 + Expect(db.Raw( + `SELECT count(*) FROM information_schema.tables WHERE table_schema = 'public'`, + ).Scan(&tables).Error).To(Succeed()) + Expect(tables).To(BeZero(), "a spec was handed a database another spec had already migrated") + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 788f8dee2d55..1ad6822afd23 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -104,15 +104,17 @@ The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_R A worker can open one long-lived, multiplexed tunnel to the frontend instead of listening on a port of its own. It dials `GET /api/cluster/connect?id=`, the connection is upgraded to a WebSocket, and every subsequent request the frontend makes to that worker travels as a stream inside it. Nothing dials *into* the worker, so a worker behind NAT, in another Kubernetes cluster or on a laptop needs no inbound port and no reachable address. -The dial is authenticated against **that node's own stored token**, not the deployment-wide registration token: the frontend hashes the presented bearer token and compares it with the hash recorded on the node's row at registration. A worker that presents a token belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity. +The dial is authenticated against **the hash stored on that node's row, which today is still the registration token's hash**: the frontend hashes the presented bearer token and compares it with what registration recorded. A worker that presents a token belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A node still awaiting admin approval is refused with `403`. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity. -One honest caveat about that today: a worker currently registers by presenting the deployment's registration token, so the hash stored on its row *is* that token's hash, and a leaked registration token plus a known node ID still gets a tunnel. What the tunnel endpoint does not do is trust the configured token directly, so when workers are issued their own secrets at registration the isolation becomes real with no change to this route. +So the isolation is not there yet, and the caveat is worth stating plainly: because a worker registers by presenting the deployment's registration token, a leaked registration token plus a known node ID still gets a tunnel. What the tunnel endpoint does not do is trust the configured token directly, so when workers are issued their own per-node secrets and those secrets are what registration stores, the isolation becomes real with no change to this route. -The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims every tunnel it still holds as soon as it re-registers. +**Worker tunnels need `LOCALAI_REGISTRATION_TOKEN` set.** A worker registering against a frontend with no registration token configured sends no token, so nothing is stored on its row, and the tunnel has nothing to authenticate it against: every dial is refused with `401`. LocalAI warns about this at startup. Setting the token later is not enough on its own; the workers have to register again for the hash to be written. + +The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims the tunnels it still holds on a live session as soon as it re-registers, skipping any whose socket has already gone. That re-claim needs the replica to have an advertised address: without one it never had an instance row to begin with, and its tunnels are usable only by the replica holding them. | Method | Path | Description | |--------|------|-------------| -| `GET` | `/api/cluster/connect?id=` | Worker opens its multiplexed tunnel (`Authorization: Bearer `) | +| `GET` | `/api/cluster/connect?id=` | Worker opens its multiplexed tunnel (`Authorization: Bearer `) | The route is exempt from the normal session/API-key authentication (it authenticates itself, like `/api/cluster/peer`) and is registered in every deployment. Outside distributed mode there is no node table to check a token against, so it answers `503`. From a816bf9b84bccaed6470f6931307f6df51eb6600 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 11:17:48 +0000 Subject: [PATCH 24/42] fix(testutil): stop the shared-database change from disarming two regressions Two advisory-lock specs named their database by literal, ALTER DATABASE testdb. Once the test helper started handing every spec its own database on a shared server, that statement landed on the maintenance database and did nothing to the one the spec was holding, so both specs went green having never reproduced the condition they exist for. They regress a model-load advisory-lock wedge that has already shipped to production once, so the previous commit's de-flaking silently disarmed a regression test for a real deployed bug. Both sites now read the name back with current_database() and, more importantly, assert the override actually landed before relying on it. A literal name can go stale again; an assertion that the setting is in force cannot pass while it is not. Removing either production override now fails the matching spec with the real 55P03 and 57014 again. That literal also meant every CREATE DATABASE and every DROP ... WITH (FORCE) ran under the 300ms bound it set on the maintenance database, which is a new load-dependent single-spec flake inside the change that was meant to remove one. The helper's maintenance connections now pin one connection and clear both timeouts on it, so no setting a spec makes can bound them, and a white-box spec imposes the leak deliberately and proves it does not reach them. Also pins the reclaimOne gate deferral the previous commit added without a test, by panicking inside the re-claim's own claim statement, and drops the per-dial empty-token log line to debug now that the boot warning says it once. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/cluster/connect.go | 7 ++- .../advisorylock/advisorylock_test.go | 61 ++++++++++++++---- core/services/cluster/tunnel_test.go | 42 +++++++++++++ core/services/testutil/testdb.go | 58 +++++++++++++++-- .../services/testutil/testdb_internal_test.go | 63 +++++++++++++++++++ 5 files changed, 214 insertions(+), 17 deletions(-) create mode 100644 core/services/testutil/testdb_internal_test.go diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go index e76450692fd9..ca94ead48126 100644 --- a/core/http/endpoints/cluster/connect.go +++ b/core/http/endpoints/cluster/connect.go @@ -111,7 +111,12 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi // secret. One log line for both leaves an operator reading "wrong token" // while every worker fails identically. if node.TokenHash == "" { - xlog.Warn("Refusing a worker tunnel: this node has no stored token, which means it registered with no registration token configured", + // Debug, not Warn. Every worker in such a deployment fails this way + // on every reconnect, so warning per dial buries the log; the fact + // is stated once, at boot, where core/http/app.go warns that no + // registration token is configured. What this line adds is which + // node, for an operator who has already read that warning. + xlog.Debug("refusing a worker tunnel: this node has no stored token, so it registered with no registration token configured", "node", nodeID, "knob", "LOCALAI_REGISTRATION_TOKEN") return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") } diff --git a/core/services/advisorylock/advisorylock_test.go b/core/services/advisorylock/advisorylock_test.go index f1bd3e75ed5c..536666f6cb5e 100644 --- a/core/services/advisorylock/advisorylock_test.go +++ b/core/services/advisorylock/advisorylock_test.go @@ -2,7 +2,9 @@ package advisorylock import ( "context" + "fmt" "runtime" + "strings" "sync" "sync/atomic" "time" @@ -14,6 +16,51 @@ import ( "gorm.io/gorm" ) +// alterThisDatabase applies a server-side setting to the database this handle is +// actually connected to, and proves it landed. +// +// The name is read back from the connection rather than written as a literal. +// The test helper hands each spec its own database on a shared server, so a +// hard-coded name ALTERs a database this handle never touches: the statement +// succeeds, the override does nothing, and the two specs below go green having +// exercised none of the condition they exist for. They regress a model-load +// advisory-lock wedge that has already shipped to production once, so a green +// spec that proves nothing is the worst outcome available here. +// +// The read-back is the guard. Idle connections are dropped first so the next one +// is opened fresh and inherits the new database-level default; SHOW then reports +// what a waiter's own connection would inherit. If that ever stops matching, the +// spec fails here rather than passing for the wrong reason. +func alterThisDatabase(db *gorm.DB, setting, value string) { + GinkgoHelper() + + var name string + Expect(db.Raw("SELECT current_database()").Scan(&name).Error).ToNot(HaveOccurred()) + Expect(name).ToNot(BeEmpty()) + + Expect(db.Exec(fmt.Sprintf("ALTER DATABASE %q SET %s = %s", name, setting, quoteLiteral(value))).Error). + ToNot(HaveOccurred()) + + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + // database/sql retains no idle connections at 0, closing the ones it is + // already holding, so every connection after this point is opened fresh and + // inherits the new database-level default. + sqlDB.SetMaxIdleConns(0) + + var applied string + Expect(db.Raw("SHOW " + setting).Scan(&applied).Error).ToNot(HaveOccurred()) + Expect(applied).To(Equal(value), + "the %s override did not reach the database this spec is holding (%s), so the spec below would pass without ever reproducing the condition it regresses", + setting, name) +} + +// quoteLiteral wraps a settings value as a SQL string literal. The values here +// are spec constants, so this only has to be correct, not hostile-input-proof. +func quoteLiteral(v string) string { + return "'" + strings.ReplaceAll(v, "'", "''") + "'" +} + var _ = Describe("AdvisoryLock", func() { Context("PostgreSQL advisory locks", func() { var db *gorm.DB @@ -166,12 +213,7 @@ var _ = Describe("AdvisoryLock", func() { // blocked on pg_advisory_lock() is aborted by the server after this // window and surfaces SQLSTATE 55P03 ("canceling statement due to // lock timeout") to the caller instead of waiting for its turn. - Expect(db.Exec("ALTER DATABASE testdb SET lock_timeout = '300ms'").Error).ToNot(HaveOccurred()) - sqlDB, err := db.DB() - Expect(err).ToNot(HaveOccurred()) - // Drop pooled connections so subsequent ones reconnect and inherit - // the new database-level lock_timeout default. - sqlDB.SetMaxIdleConns(0) + alterThisDatabase(db, "lock_timeout", "300ms") holding := make(chan struct{}) released := make(chan struct{}) @@ -214,12 +256,7 @@ var _ = Describe("AdvisoryLock", func() { // statement_timeout=60s; a cold model load holds the lock far longer, // so every concurrent caller died with SQLSTATE 57014 ("canceling // statement due to statement timeout") rather than waiting its turn. - Expect(db.Exec("ALTER DATABASE testdb SET statement_timeout = '300ms'").Error).ToNot(HaveOccurred()) - sqlDB, err := db.DB() - Expect(err).ToNot(HaveOccurred()) - // Drop pooled connections so subsequent ones reconnect and inherit - // the new database-level statement_timeout default. - sqlDB.SetMaxIdleConns(0) + alterThisDatabase(db, "statement_timeout", "300ms") holding := make(chan struct{}) released := make(chan struct{}) diff --git a/core/services/cluster/tunnel_test.go b/core/services/cluster/tunnel_test.go index acc81d81813f..7d0136b47e67 100644 --- a/core/services/cluster/tunnel_test.go +++ b/core/services/cluster/tunnel_test.go @@ -715,6 +715,48 @@ var _ = Describe("Re-claiming tunnels after this replica's rows were reaped", fu } Consistently(stored, 2*cluster.InstanceHeartbeat, time.Second).Should(Equal(epoch)) }) + + It("frees the node's gate when a re-claim panics", func() { + // The same property Attach's gate specs pin, on the other function that + // takes the gate. Reclaim runs from the heartbeat loop, so a wedged gate + // here is worse than one wedged by a dial: nothing retries it, and the + // worker can never re-attach to this replica because its Attach blocks + // in enterClaim until its own context expires. + // + // The panic is thrown from inside the re-claim's own Claim statement, + // on that statement's goroutine, using the same gorm Trace hook the + // interleaving specs above use. That is the window a real panic under + // Claim would land in: the row is written and the gate is held. + hook := newClaimHook(isClaimOf("w1")) + hooked := cluster.NewTunnelRegistry( + cluster.NewRegistry(db.Session(&gorm.Session{Logger: hook})), "me") + + frontend, _ := workerTunnel() + _, err := hooked.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + // Installed after the attach so the hook fires on the RE-claim, not on + // the claim that set the tunnel up. + hook.setAction(func(string) { panic("claim exploded") }) + + panicked := func() (p bool) { + defer func() { p = recover() != nil }() + _, _ = hooked.Reclaim(ctx) + return + }() + Expect(panicked).To(BeTrue(), + "the hook did not fire inside the re-claim, so this spec is no longer testing what it claims") + + // A wedged gate is indistinguishable from a slow one except by waiting. + // The hook has already fired once and will not fire again, so this + // Attach either completes or never reaches Claim at all. + bounded, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + next, _ := workerTunnel() + _, err = hooked.Attach(bounded, "w1", next) + Expect(err).ToNot(HaveOccurred(), + "the panicking re-claim left this node's gate closed, so the worker can never attach to this replica again") + }) }) var _ = Describe("The worker tunnel registry's claim gate", func() { diff --git a/core/services/testutil/testdb.go b/core/services/testutil/testdb.go index 755589cd438c..7e2c061dd8a4 100644 --- a/core/services/testutil/testdb.go +++ b/core/services/testutil/testdb.go @@ -146,7 +146,7 @@ func SetupTestDB() *gorm.DB { // teardown. closePool(db) - drop, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) + drop, err := openTolerantPool(dsn) if err != nil { // Reported, never asserted. A cleanup that fails the spec turns one // database hiccup into a failure that buries whatever the spec was @@ -165,16 +165,66 @@ func SetupTestDB() *gorm.DB { return db } -// openPool connects to dsn with logging off. Used for the short-lived -// maintenance connections only; the database a spec is handed keeps gorm's -// silent logger so a caller can still swap it. +// openPool connects to dsn with logging off and with every server-side timeout +// disabled on the session. Used for the short-lived maintenance connections +// only; the database a spec is handed keeps gorm's silent logger and the +// server's defaults, because setting timeouts on it is a thing specs do on +// purpose. +// +// The timeouts are cleared because CREATE DATABASE and DROP DATABASE must not +// be bounded by anything a spec configured. A spec that sets a short +// statement_timeout on ITS database cannot reach this one, but a spec that +// names the maintenance database by mistake can, and that is not hypothetical: +// two advisory-lock specs did exactly that until this round. The consequence +// there is a load-dependent failure in another spec's setup or a silently +// swallowed DROP, which is the same invisible single-spec flake this helper was +// rewritten to remove. +// +// MaxOpenConns(1) is what makes the SET reach the statement that follows it: a +// session setting lives on one connection, and with a single connection in the +// pool there is no other one for CREATE or DROP to land on. func openPool(dsn string) *gorm.DB { GinkgoHelper() db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) Expect(err).ToNot(HaveOccurred()) + + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + sqlDB.SetMaxOpenConns(1) + + Expect(db.Exec("SET statement_timeout = 0").Error).To(Succeed()) + Expect(db.Exec("SET lock_timeout = 0").Error).To(Succeed()) return db } +// openTolerantPool is openPool for the cleanup path, which must report a +// failure rather than assert one: an assertion here would fail a spec that had +// already passed, and bury whatever the next real failure was. +func openTolerantPool(dsn string) (*gorm.DB, error) { + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) + if err != nil { + return nil, err + } + sqlDB, err := db.DB() + if err != nil { + closePool(db) + return nil, err + } + sqlDB.SetMaxOpenConns(1) + // The DROP below is the statement most likely to be slow, since it waits on + // FORCE terminating other sessions, so it is the one a leaked timeout would + // abort. Measured at up to 169ms under load, against a 300ms bound. + if err := db.Exec("SET statement_timeout = 0").Error; err != nil { + closePool(db) + return nil, err + } + if err := db.Exec("SET lock_timeout = 0").Error; err != nil { + closePool(db) + return nil, err + } + return db, nil +} + func closePool(db *gorm.DB) { if db == nil { return diff --git a/core/services/testutil/testdb_internal_test.go b/core/services/testutil/testdb_internal_test.go new file mode 100644 index 000000000000..457d6f7d5f5f --- /dev/null +++ b/core/services/testutil/testdb_internal_test.go @@ -0,0 +1,63 @@ +package testutil + +import ( + "fmt" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// These are white-box on purpose: the property is about the connection this +// package makes for itself, which no caller can reach. +var _ = Describe("the maintenance connection", func() { + It("cannot be bounded by a timeout set on the maintenance database", func() { + // The leak this pins is not hypothetical. Two advisory-lock specs named + // a database by literal, and once the helper started handing out + // per-spec databases those ALTERs landed on the maintenance database + // instead, so every CREATE DATABASE and every DROP ... WITH (FORCE) ran + // under a 300ms bound. A CREATE that trips it fails another spec's + // setup; a DROP that trips it is swallowed and leaks a database. Both + // are load-dependent single-spec failures, which is the exact shape + // this helper was rewritten to remove. + dsn := sharedPostgres() + + var maintenance string + func() { + probe := openPool(dsn) + defer closePool(probe) + Expect(probe.Raw("SELECT current_database()").Scan(&maintenance).Error).To(Succeed()) + }() + Expect(maintenance).ToNot(BeEmpty()) + + // Impose the leak, then assert a fresh maintenance connection is + // unaffected. Reset first so a failure below cannot leave the bound in + // place for the rest of the suite. + DeferCleanup(func() { + reset := openPool(dsn) + defer closePool(reset) + Expect(reset.Exec(fmt.Sprintf("ALTER DATABASE %q RESET statement_timeout", maintenance)).Error).To(Succeed()) + Expect(reset.Exec(fmt.Sprintf("ALTER DATABASE %q RESET lock_timeout", maintenance)).Error).To(Succeed()) + }) + func() { + impose := openPool(dsn) + defer closePool(impose) + Expect(impose.Exec(fmt.Sprintf("ALTER DATABASE %q SET statement_timeout = '1ms'", maintenance)).Error).To(Succeed()) + Expect(impose.Exec(fmt.Sprintf("ALTER DATABASE %q SET lock_timeout = '1ms'", maintenance)).Error).To(Succeed()) + }() + + fresh := openPool(dsn) + defer closePool(fresh) + var statementTimeout, lockTimeout string + Expect(fresh.Raw("SHOW statement_timeout").Scan(&statementTimeout).Error).To(Succeed()) + Expect(fresh.Raw("SHOW lock_timeout").Scan(&lockTimeout).Error).To(Succeed()) + Expect(statementTimeout).To(Equal("0"), + "a statement_timeout on the maintenance database reached the helper's own connection, so CREATE and DROP DATABASE are bounded by whatever a spec configured") + Expect(lockTimeout).To(Equal("0"), + "a lock_timeout on the maintenance database reached the helper's own connection") + + // And the thing the timeouts would actually abort still works while the + // bound is in force. A 1ms statement_timeout is far below the 14-26ms a + // CREATE DATABASE takes here, so this could not pass by being fast. + Expect(SetupTestDB()).ToNot(BeNil()) + }) +}) From 5b7d65e67940d8957eb8b319ded5bfcda99da144 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 11:43:23 +0000 Subject: [PATCH 25/42] fix(testutil): clear the maintenance timeouts at connection startup The guard added last commit was circular. It cleared the maintenance database's timeouts by executing SET statement_timeout = 0 on a connection that had already inherited that database's bound, so the statement clearing the bound ran under the bound it was clearing. Under the white-box spec's deliberate 1ms that gave it 1ms, and it failed roughly once in fifty at 8-way concurrency with SQLSTATE 57014. The guard against invisible load-dependent flakes had become one. The clearing is now delivered as a connection startup option, options=-c statement_timeout=0 -c lock_timeout=0 on the maintenance DSN, so there is no statement left to abort. Raising the imposed bound would only have bought headroom and left the circularity in place. pgx puts every URL query parameter into settings, options is absent from notRuntimeParams so it becomes a runtime parameter, and runtime parameters are copied into the startup message (pgconn/config.go:340-378, 606-617; pgconn/pgconn.go:382-388). The spec now discriminates on pg_settings.reset_val, the value in force when the connection started: 0 for a startup option, 1ms for a session SET. A first attempt using a deliberately slow first statement did NOT discriminate, because under the circular design the SET is itself the first statement, so by the time a spec runs anything the session is already unbounded. Reinstating the circular clearing now reddens the spec deterministically rather than intermittently. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/testutil/testdb.go | 87 ++++++++++--------- .../services/testutil/testdb_internal_test.go | 53 ++++++++++- 2 files changed, 98 insertions(+), 42 deletions(-) diff --git a/core/services/testutil/testdb.go b/core/services/testutil/testdb.go index 7e2c061dd8a4..4f36d52c3d31 100644 --- a/core/services/testutil/testdb.go +++ b/core/services/testutil/testdb.go @@ -165,61 +165,70 @@ func SetupTestDB() *gorm.DB { return db } -// openPool connects to dsn with logging off and with every server-side timeout -// disabled on the session. Used for the short-lived maintenance connections -// only; the database a spec is handed keeps gorm's silent logger and the -// server's defaults, because setting timeouts on it is a thing specs do on -// purpose. +// maintenanceDSN is dsn with every server-side timeout disabled as a CONNECTION +// STARTUP OPTION rather than as a statement. // -// The timeouts are cleared because CREATE DATABASE and DROP DATABASE must not -// be bounded by anything a spec configured. A spec that sets a short -// statement_timeout on ITS database cannot reach this one, but a spec that -// names the maintenance database by mistake can, and that is not hypothetical: -// two advisory-lock specs did exactly that until this round. The consequence -// there is a load-dependent failure in another spec's setup or a silently -// swallowed DROP, which is the same invisible single-spec flake this helper was -// rewritten to remove. +// The timeouts have to go because CREATE DATABASE and DROP DATABASE must not be +// bounded by anything a spec configured. A spec that sets a short +// statement_timeout on ITS own database cannot reach this connection, but a spec +// that names the maintenance database by mistake can, and that is not +// hypothetical: two advisory-lock specs did exactly that. // -// MaxOpenConns(1) is what makes the SET reach the statement that follows it: a -// session setting lives on one connection, and with a single connection in the -// pool there is no other one for CREATE or DROP to land on. +// Clearing it with `SET statement_timeout = 0` on an already-open connection is +// circular and was a real defect here: that connection has already inherited the +// database's bound, so the statement that clears the bound runs under it and can +// be aborted by it with SQLSTATE 57014. It failed roughly once in fifty at +// 8-way concurrency, which is the same invisible load-dependent single-spec +// flake this helper exists to remove. A startup option removes the circularity +// instead of buying headroom against it: the value is delivered in the startup +// packet, so the connection is already unbounded before it can run anything. +// +// The route is verified in the driver rather than assumed. pgx puts every URL +// query parameter into settings (pgconn/config.go:614), `options` is absent from +// notRuntimeParams (pgconn/config.go:340-362) so it becomes a runtime parameter +// (pgconn/config.go:374-378), and runtime parameters are copied into the startup +// message (pgconn/pgconn.go:382-388). PostgreSQL treats `options` as backend +// command-line switches, so `-c statement_timeout=0` is applied before the +// session accepts a query. +func maintenanceDSN(dsn string) (string, error) { + u, err := url.Parse(dsn) + if err != nil { + return "", err + } + q := u.Query() + // Percent-encoded by Encode, and pgx decodes query values before they reach + // settings, so the server receives the switches with their spaces intact. + q.Set("options", "-c statement_timeout=0 -c lock_timeout=0") + u.RawQuery = q.Encode() + return u.String(), nil +} + +// openPool connects to the maintenance database with logging off and no +// server-side timeouts. Used for the short-lived maintenance connections only; +// the database a spec is handed keeps gorm's silent logger and the server's +// defaults, because setting timeouts on it is a thing specs do on purpose. func openPool(dsn string) *gorm.DB { GinkgoHelper() - db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) - Expect(err).ToNot(HaveOccurred()) - - sqlDB, err := db.DB() + db, err := openTolerantPool(dsn) Expect(err).ToNot(HaveOccurred()) - sqlDB.SetMaxOpenConns(1) - - Expect(db.Exec("SET statement_timeout = 0").Error).To(Succeed()) - Expect(db.Exec("SET lock_timeout = 0").Error).To(Succeed()) return db } // openTolerantPool is openPool for the cleanup path, which must report a // failure rather than assert one: an assertion here would fail a spec that had // already passed, and bury whatever the next real failure was. +// +// It carries the same startup options, and the DROP is the statement that most +// needs them: FORCE waits on terminating other sessions, measured at up to 169ms +// against the 300ms bound that used to leak here, and a DROP aborted mid-way is +// swallowed and leaks a database. func openTolerantPool(dsn string) (*gorm.DB, error) { - db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard}) + maintenance, err := maintenanceDSN(dsn) if err != nil { return nil, err } - sqlDB, err := db.DB() + db, err := gorm.Open(postgres.Open(maintenance), &gorm.Config{Logger: logger.Discard}) if err != nil { - closePool(db) - return nil, err - } - sqlDB.SetMaxOpenConns(1) - // The DROP below is the statement most likely to be slow, since it waits on - // FORCE terminating other sessions, so it is the one a leaked timeout would - // abort. Measured at up to 169ms under load, against a 300ms bound. - if err := db.Exec("SET statement_timeout = 0").Error; err != nil { - closePool(db) - return nil, err - } - if err := db.Exec("SET lock_timeout = 0").Error; err != nil { - closePool(db) return nil, err } return db, nil diff --git a/core/services/testutil/testdb_internal_test.go b/core/services/testutil/testdb_internal_test.go index 457d6f7d5f5f..b3003c0d0016 100644 --- a/core/services/testutil/testdb_internal_test.go +++ b/core/services/testutil/testdb_internal_test.go @@ -5,6 +5,9 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "gorm.io/driver/postgres" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" ) // These are white-box on purpose: the property is about the connection this @@ -45,6 +48,11 @@ var _ = Describe("the maintenance connection", func() { Expect(impose.Exec(fmt.Sprintf("ALTER DATABASE %q SET lock_timeout = '1ms'", maintenance)).Error).To(Succeed()) }() + // The bound is delivered before the first statement, so the check + // below is also the connection's first statement. That ordering is the + // point: clearing the bound with a SET would be circular, because the + // clearing statement inherits the bound it is clearing and can be + // aborted by it with 57014. There is no such bootstrap statement now. fresh := openPool(dsn) defer closePool(fresh) var statementTimeout, lockTimeout string @@ -55,9 +63,48 @@ var _ = Describe("the maintenance connection", func() { Expect(lockTimeout).To(Equal("0"), "a lock_timeout on the maintenance database reached the helper's own connection") - // And the thing the timeouts would actually abort still works while the - // bound is in force. A 1ms statement_timeout is far below the 14-26ms a - // CREATE DATABASE takes here, so this could not pass by being fast. + // A control, and the reason this spec is not a race. Clearing the bound + // with a statement is circular: the clearing statement runs on a + // connection that has already inherited the bound. Whether that + // particular statement exceeds 1ms is a matter of load, which makes the + // defect an intermittent one; whether the FIRST statement on a plain + // connection is bounded at all is not. So the control asks the + // deterministic question, with a first statement that certainly exceeds + // the bound. + func() { + plain, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: gormlogger.Discard}) + Expect(err).ToNot(HaveOccurred()) + defer closePool(plain) + err = plain.Exec("SELECT pg_sleep(0.05)").Error + Expect(err).To(HaveOccurred(), + "the imposed bound does not reach a fresh connection's first statement, so this spec's subject is not actually under test") + Expect(err.Error()).To(ContainSubstring("57014"), + "expected the imposed statement_timeout to abort this, got something else") + }() + + // The same first statement on a maintenance connection is unbounded. + Expect(fresh.Exec("SELECT pg_sleep(0.05)").Error).To(Succeed()) + + // And this is the assertion that says WHY, which is the part a + // statement-based clearing cannot satisfy. reset_val is the value the + // session would fall back to, that is, the value that was in force when + // the connection started, before it could run anything. Clearing the + // bound with `SET statement_timeout = 0` leaves reset_val at the + // database's 1ms: the session is unbounded only because a statement + // said so, and that statement ran under the 1ms bound and can be + // aborted by it. Delivering it as a startup option makes the connection + // unbounded with no statement in between, which is the difference + // between a fix and a wider margin. + var resetVal string + Expect(fresh.Raw( + "SELECT reset_val FROM pg_settings WHERE name = 'statement_timeout'", + ).Scan(&resetVal).Error).To(Succeed()) + Expect(resetVal).To(Equal("0"), + "the maintenance connection started under a %s bound and cleared it with a statement, so the clearing statement itself runs under the bound it is clearing", resetVal) + + // And the operation the bound would abort still works while it is in + // force. 1ms is far below the 14-26ms a CREATE DATABASE takes here, so + // this cannot pass by being fast. Expect(SetupTestDB()).ToNot(BeNil()) }) }) From 29a2020f3d0a3a71695535b2883b25f59986fa36 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 12:28:00 +0000 Subject: [PATCH 26/42] feat(worker): dial, hold and serve the tunnel, on a credential of its own The worker end of the tunnel. It dials wss:///api/cluster/connect, holds one yamux session as the CLIENT, and serves every stream the frontend opens on it. Nothing dials into the worker, which is the point: no inbound port, no reachable address. Each stream opens with a length-prefixed frame naming a tag and a target, and the worker answers before either side speaks the tunnelled protocol. The reply is sent on every stream, not only on refusal, because the protocols carried here are client-speaks-first and a reply sent only sometimes would arrive interleaved with a response body. Two tags today: grpc reaches a backend process, and only on 127.0.0.1 within this worker's own backend port range, because a tunnel terminates inside the worker and letting the frontend name a host would make every worker a proxy into its own LAN; http reaches the worker's file-transfer server, whose address the frontend is not asked about. An unknown tag, an unreachable local service and an unparseable request are three refusals and stay three on the wire. A frontend gives up on the first and retries the second. Each is answered AND the stream is ended: a worker that says why and leaves the stream open has parked the caller on a request nobody will answer, and a deadline on the far side cannot tell that from a slow worker. The specs assert the stream ends rather than that an error occurred, which is what phase 1 shipped in three places and held in none. Reconnects double from 500ms to a 30s ceiling, each wait drawn between half the interval and all of it, and the interval returns to its floor only after a session that LASTED. Resetting on connect is how a rolling restart, where every dial succeeds and dies moments later, becomes a retry storm against the first replica back up. Nothing is assumed to survive a reconnect: the credential is read at dial time, never captured. And the credential is now real. The tunnel endpoint advertised authenticating a worker against its own secret, but registration stored the hash of the shared registration token, so a leak plus a known node ID still opened a tunnel. Registration now mints a per-node secret, returns the plaintext once as tunnel_token, and stores only its SHA-256 in a new column; the endpoint compares against that and does not fall back to the old one. Rotating on every registration follows from storing only the hash, since a re-registering worker cannot be told the secret it already holds; its live tunnel is unaffected, because the credential is checked when a tunnel is dialled and never again. Unlike the agent API key and the NATS JWT next to it, the credential IS issued to a node awaiting approval: the tunnel route re-reads the node's status on every dial and refuses a pending one, so it is inert until an admin acts, and withholding it would strand every worker that registers exactly once. A node that has not registered since this change cannot tunnel, and the column cannot be back-filled because the plaintext only ever existed in the response that minted it. The boot warning that said tunnels need LOCALAI_REGISTRATION_TOKEN is replaced: it was true while the tunnel authenticated against that token's hash, and says the wrong thing now. What is still true, and is what it warns about instead, is that without one, registration itself is unauthenticated. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/cli/workerregistry/client.go | 37 +- core/cli/workerregistry/credentials.go | 21 + core/http/app.go | 25 +- core/http/endpoints/cluster/connect.go | 62 +- core/http/endpoints/cluster/connect_test.go | 53 +- core/http/endpoints/localai/nodes.go | 39 ++ core/http/endpoints/localai/nodes_test.go | 88 +++ core/services/cluster/tunnelproto.go | 239 ++++++++ core/services/nodes/registry.go | 48 +- core/services/worker/config.go | 7 +- core/services/worker/tunnel.go | 617 ++++++++++++++++++++ core/services/worker/tunnel_test.go | 509 ++++++++++++++++ core/services/worker/worker.go | 59 +- docs/content/features/distributed-mode.md | 32 +- 14 files changed, 1748 insertions(+), 88 deletions(-) create mode 100644 core/services/cluster/tunnelproto.go create mode 100644 core/services/worker/tunnel.go create mode 100644 core/services/worker/tunnel_test.go diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index cf46455c95c0..f8728166fef9 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -58,9 +58,15 @@ func (c *RegistrationClient) setAuth(req *http.Request) { // RegisterResponse is the JSON body returned by /api/node/register. type RegisterResponse struct { - ID string `json:"id"` - Status string `json:"status,omitempty"` // "pending" until an admin approves the node - APIToken string `json:"api_token,omitempty"` + ID string `json:"id"` + Status string `json:"status,omitempty"` // "pending" until an admin approves the node + APIToken string `json:"api_token,omitempty"` + // TunnelToken is this node's own credential for GET /api/cluster/connect. + // The frontend mints a fresh one on every registration and keeps only its + // hash, so this is the ONLY time the plaintext exists anywhere but in this + // worker's memory: a worker that discards it cannot get it back without + // registering again. + TunnelToken string `json:"tunnel_token,omitempty"` NatsJWT string `json:"nats_jwt,omitempty"` NatsUserSeed string `json:"nats_user_seed,omitempty"` } @@ -108,27 +114,42 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) } // RegisterWithRetry retries registration with exponential backoff. +// +// It drops every field of the response it does not name, the tunnel credential +// among them. Callers that need one use RegisterFullWithRetry. 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 retries registration with exponential backoff and +// returns the whole response. +func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map[string]any, maxRetries int) (*RegisterResponse, error) { 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..f9d8c2231ddc 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -50,6 +50,11 @@ type NATSCredentialManager struct { jwt string seed string nodeID string + // tunnelToken is the node's own tunnel credential from the most recent + // registration. It is kept here because every re-registration this manager + // performs ROTATES it, so the tunnel client has to read the current value + // at dial time rather than be handed one at startup. + tunnelToken string } // NewNATSCredentialManager builds a manager over register. When requireCreds is @@ -87,6 +92,22 @@ func (m *NATSCredentialManager) store(res *RegisterResponse) { if res.NatsJWT != "" && res.NatsUserSeed != "" { m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed } + // Guarded the same way the NATS pair is: a response that carries no tunnel + // token (a frontend that predates them, or one whose minting failed) must + // not wipe a working credential this worker already holds. Overwriting with + // "" would lock the tunnel out until the next registration that did carry + // one, which is the opposite of what an empty field means. + if res.TunnelToken != "" { + m.tunnelToken = res.TunnelToken + } +} + +// TunnelToken returns the node's current tunnel credential, empty until one has +// been issued. It is the callback the tunnel client reads on every dial. +func (m *NATSCredentialManager) TunnelToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.tunnelToken } // Current returns the latest NATS credentials (both empty until acquired). diff --git a/core/http/app.go b/core/http/app.go index 8f4f34912d24..fe790fb16e4e 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -603,16 +603,21 @@ func API(application *application.Application) (*echo.Echo, error) { if d := application.Distributed(); d != nil { tunnels = d.Tunnels if distCfg.RegistrationToken == "" { - // A separate warning from the peer link's, for the same missing - // knob, because the broken thing is different and an operator has - // to be told both. A worker registering against a frontend with no - // registration token sends no token, so RegisterNodeEndpoint stores - // an empty token_hash, and the tunnel authenticates against exactly - // that column: every dial 401s, forever, on a deployment that looks - // correctly configured. Fixing it needs the token set AND the - // workers re-registered, which is why it is worth saying at boot - // rather than leaving in the 401s. - xlog.Warn("Worker tunnels will refuse every dial: no registration token is configured, so no worker has a stored token to authenticate against", + // A different warning from the peer link's, for the same missing + // knob, because what breaks is different. Tunnels themselves work + // without a registration token: each node is minted its own tunnel + // credential at registration whether or not one is configured. What + // is missing is the gate in FRONT of that. With no registration + // token, RegisterNodeEndpoint validates nothing, so anyone who can + // reach this frontend can register a node and be handed a working + // tunnel credential for it. + // + // This warning replaced one that said the opposite, that tunnels + // would refuse every dial without this token. That was true while + // the tunnel authenticated against the registration token's own + // hash, and stopped being true when nodes got credentials of their + // own. + xlog.Warn("Node registration is unauthenticated, so any caller that can reach this frontend can register a worker and be issued a tunnel credential", "route", clustersvc.ConnectPath, "knob", "LOCALAI_REGISTRATION_TOKEN") } } diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go index ca94ead48126..5e6ace97f7b6 100644 --- a/core/http/endpoints/cluster/connect.go +++ b/core/http/endpoints/cluster/connect.go @@ -104,23 +104,21 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi } // Split from the mismatch below because they are different operator - // problems with different fixes. An empty stored hash means the node - // registered against a frontend with no LOCALAI_REGISTRATION_TOKEN, so - // no worker on that deployment can ever tunnel until one is set and the - // workers re-register; a mismatch means one worker holds the wrong - // secret. One log line for both leaves an operator reading "wrong token" - // while every worker fails identically. - if node.TokenHash == "" { - // Debug, not Warn. Every worker in such a deployment fails this way - // on every reconnect, so warning per dial buries the log; the fact - // is stated once, at boot, where core/http/app.go warns that no - // registration token is configured. What this line adds is which - // node, for an operator who has already read that warning. - xlog.Debug("refusing a worker tunnel: this node has no stored token, so it registered with no registration token configured", - "node", nodeID, "knob", "LOCALAI_REGISTRATION_TOKEN") + // problems with different fixes. An empty stored hash means this node + // last registered against a LocalAI that predates per-node tunnel + // credentials, so it holds no secret this route can check and must + // register again; a mismatch means the worker is presenting the wrong + // one, usually a stale credential from before a rotation. One log line + // for both leaves an operator reading "wrong token" while a whole fleet + // of not-yet-restarted workers fails identically. + if node.TunnelTokenHash == "" { + // Debug, not Warn. Every worker in that state fails this way on + // every reconnect, so warning per dial buries the log. + 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.TokenHash) { + if !authorizedWorker(token, node.TunnelTokenHash) { xlog.Debug("worker tunnel dial presented the wrong token", "node", nodeID) return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") } @@ -229,32 +227,26 @@ func bearerToken(r *http.Request) (string, bool) { // authorizedWorker compares a presented token against the hash stored on the // node's own row, in constant time. // -// Against the NODE's stored hash, not the configured registration token, which -// is a deliberate strengthening over how /api/node/ authenticates. A tunnel is a -// durable, multiplexed pipe into a worker; a credential that authorizes every -// worker at once would mean one leak lets an attacker impersonate any worker -// whose ID it can read, and take over that worker's traffic by claiming its -// tunnel. +// Against the NODE's OWN tunnel credential, not the deployment's registration +// token. A tunnel is a durable, multiplexed pipe into a worker, and a +// credential that authorizes every worker at once would mean one leak lets an +// attacker impersonate any worker whose ID it can read and take over that +// worker's traffic by claiming its tunnel. The secret compared here is minted +// per node at registration (attachTunnelToken in +// core/http/endpoints/localai/nodes.go), returned to that worker once, and +// stored only as this hash, so knowing the registration token no longer gets +// anyone a tunnel. // -// What that buys TODAY is the mechanism, not yet the isolation, and the -// difference should not be overstated: a worker registers by presenting the -// deployment's registration token, so the hash on its row is that token's hash -// and a leaked registration token still opens a tunnel for a node whose ID the -// attacker knows. What this does rule out is the shortcut of comparing against -// the configured token itself, which is what would have to be unpicked later; -// the day a worker is issued its own secret at registration, this check starts -// isolating workers from each other with no change here, PROVIDED the secret -// lands in TokenHash. The one per-node secret LocalAI mints today, the agent -// worker's api_token, does not: provisionAgentWorkerKey writes an auth.User and -// an auth.APIKey referenced by node.AuthUserID / node.APIKeyID and never touches -// this column. If the next task follows that precedent instead, this comparison -// is what has to change. +// Note which column: BackendNode.TunnelTokenHash, not TokenHash. TokenHash is +// still the hash of whatever token the worker registered WITH, which on most +// deployments is the shared registration token, and comparing against it is +// exactly the weakness this replaced. // // The empty-hash guard is defensive rather than deciding: a stored hash is // hex-encoded SHA-256, so 64 bytes or nothing, and ConstantTimeCompare already // returns 0 on a length mismatch (crypto/internal/fips140/subtle/constant_time.go:17-20 // returns 0 outright when the lengths differ). It is kept because a reader should not have to -// derive "an unregistered node authorizes nobody" from a length rule, and +// derive "a node with no credential authorizes nobody" from a length rule, and // because the caller logs that case separately. func authorizedWorker(token, storedHash string) bool { if storedHash == "" { diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go index 072f69d0226c..15206d50444c 100644 --- a/core/http/endpoints/cluster/connect_test.go +++ b/core/http/endpoints/cluster/connect_test.go @@ -24,12 +24,17 @@ import ( "gorm.io/gorm" ) -// workerToken is the secret one worker holds. It is stored on that worker's row -// as a hash, never in a config value, which is the whole point of the check the -// specs below pin: a second worker's token, and the deployment-wide +// workerToken is the tunnel credential one worker holds. Registration mints it +// per node and keeps only its hash, which is the whole point of the check the +// specs below pin: a second worker's credential, and the deployment-wide // registration token, are both wrong for this node. const workerToken = "worker-1-secret" +// registrationToken stands in for the shared secret every worker in a +// deployment registers with. It is stored on the row too, in a DIFFERENT +// column, and a spec below pins that presenting it does not open a tunnel. +const registrationToken = "deployment-registration-token" + func tokenHash(token string) string { sum := sha256.Sum256([]byte(token)) return hex.EncodeToString(sum[:]) @@ -66,9 +71,13 @@ var _ = Describe("Worker tunnel handler", func() { Expect(err).ToNot(HaveOccurred()) node := &nodes.BackendNode{ - Name: "worker-1", - Address: "10.0.0.9:50051", - TokenHash: tokenHash(workerToken), + Name: "worker-1", + Address: "10.0.0.9:50051", + // Both hashes are set, and they differ. That is what a real + // registration produces: TokenHash is the shared token the worker + // registered WITH, TunnelTokenHash is the secret minted FOR it. + TokenHash: tokenHash(registrationToken), + TunnelTokenHash: tokenHash(workerToken), } Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) nodeID = node.ID @@ -108,11 +117,13 @@ var _ = Describe("Worker tunnel handler", func() { Expect(tun.Held()).To(BeEmpty()) }) - It("refuses a token that is not the node's own", func() { + It("refuses the deployment's registration token, which this node also stores", func() { // The registration token is the credential every worker in the - // deployment holds. Accepting it here would mean one leaked shared - // secret impersonates any worker whose ID an attacker can read. - _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer("deployment-registration-token")) + // deployment holds, and it IS on this node's row, in TokenHash. + // Accepting it would mean one leaked shared secret impersonates any + // worker whose ID an attacker can read; the node's own tunnel + // credential is the only thing this route accepts. + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(registrationToken)) Expect(err).To(HaveOccurred()) Expect(resp).ToNot(BeNil()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) @@ -194,13 +205,19 @@ var _ = Describe("Worker tunnel handler", func() { "the claim outlived the socket, so this replica keeps being named the owner of a worker it no longer holds") }) - It("refuses a node whose row carries no stored token", func() { - // A deployment with no registration token configured produces exactly - // these rows: the worker sends no token, so registration stores none, - // and there is nothing here to authenticate against. - Expect(db.Exec(`UPDATE backend_nodes SET token_hash = '' WHERE id = ?`, nodeID).Error).To(Succeed()) - - _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken)) + It("refuses a node with no tunnel credential, without falling back to its registration token", func() { + // A node registered by a LocalAI predating per-node tunnel credentials + // produces exactly this row: a registration-token hash in token_hash + // and nothing in tunnel_token_hash. It cannot be back-filled, because + // the plaintext only ever existed in the response that minted it, so + // such a node must register again. + // + // The dial presents the REGISTRATION token, which is still on the row. + // A handler that fell back to token_hash when the tunnel hash is empty + // would let it in, which is the weakness this whole change removed. + Expect(db.Exec(`UPDATE backend_nodes SET tunnel_token_hash = '' WHERE id = ?`, nodeID).Error).To(Succeed()) + + _, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(registrationToken)) Expect(err).To(HaveOccurred()) Expect(resp).ToNot(BeNil()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) @@ -266,7 +283,7 @@ var _ = Describe("Worker tunnel handler when the attach panics", func() { db = testutil.SetupTestDB() nodeReg, err := nodes.NewNodeRegistry(db) Expect(err).ToNot(HaveOccurred()) - node := &nodes.BackendNode{Name: "worker-1", Address: "10.0.0.9:50051", TokenHash: tokenHash(workerToken)} + node := &nodes.BackendNode{Name: "worker-1", Address: "10.0.0.9:50051", TunnelTokenHash: tokenHash(workerToken)} Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) nodeID = node.ID diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index bbae523b1025..0be399a06f49 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" @@ -244,6 +245,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } + attachTunnelToken(ctx, response, registry, node) attachNatsJWT(response, node, natsCfg) return c.JSON(http.StatusCreated, response) @@ -288,6 +290,43 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr } } +// attachTunnelToken mints this node a fresh tunnel credential, stores only its +// hash, and puts the plaintext in the registration response. +// +// It is minted for EVERY node that registers, pending ones included, which is a +// deliberate divergence from the two other per-node credentials in this file: +// the agent worker's API key (provisionAgentWorkerKey) and its NATS JWT +// (attachNatsJWT) are both withheld from a node awaiting approval. Those two +// are bearer grants that WORK the moment they are issued, so issuing one to an +// unapproved node would route around the admin. A tunnel credential is not: +// core/http/endpoints/cluster/connect.go re-reads the node's status on every +// dial and refuses a pending node with 403, so the credential is inert until an +// admin approves and stays inert if approval is revoked. Withholding it would +// instead strand every worker that registers exactly once (the static-NATS +// path in core/services/worker/worker.go does), because approval alone does not +// prompt a re-registration and nothing else can hand it the secret. +// +// A failure to mint or to store is logged and the response goes out without the +// token. Registration is what gets a worker into the cluster at all, and +// failing it over a credential the worker does not need until it tunnels would +// turn a tunnel problem into a node that cannot join. The worker sees no +// tunnel_token, reports that it has no credential, and retries at its next +// registration. +func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) { + if node == nil { + return + } + // crypto/rand.Text: at least 128 bits of randomness, no error to handle and + // no length constant to get wrong. + plaintext := rand.Text() + sum := sha256.Sum256([]byte(plaintext)) + 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"] = plaintext +} + // 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_test.go b/core/http/endpoints/localai/nodes_test.go index 19e6a6b07eea..5b2d6841462e 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "crypto/subtle" + "encoding/hex" "encoding/json" "net/http" "net/http/httptest" @@ -20,6 +21,12 @@ import ( . "github.com/onsi/gomega" ) +// hashOf is how the node registry stores a secret: hex-encoded SHA-256. +func hashOf(secret string) string { + sum := sha256.Sum256([]byte(secret)) + return hex.EncodeToString(sum[:]) +} + var _ = DescribeTable("token validation", func(expectedToken, providedToken string, wantMatch bool) { if expectedToken == "" { @@ -77,6 +84,87 @@ var _ = Describe("Node HTTP handlers", func() { Expect(resp["status"]).To(Equal(nodes.StatusHealthy)) }) + // register posts one registration and returns the decoded response. + register := func(body string, expectedToken string, autoApprove bool) map[string]any { + 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) + + handler := RegisterNodeEndpoint(registry, expectedToken, autoApprove, nil, "", natsauth.Config{}) + ExpectWithOffset(1, handler(c)).To(Succeed()) + ExpectWithOffset(1, rec.Code).To(Equal(http.StatusCreated)) + + var resp map[string]any + ExpectWithOffset(1, json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp + } + + It("mints a per-node tunnel credential and stores only its hash", func() { + resp := register(`{"name":"worker-tunnel","address":"10.0.0.3:50051","token":"shared-registration-token"}`, + "shared-registration-token", true) + + plaintext, _ := resp["tunnel_token"].(string) + Expect(plaintext).ToNot(BeEmpty()) + // Not the registration token. That is the whole point: a leaked + // registration token plus a known node ID used to open a tunnel, + // because the tunnel authenticated against the hash of exactly the + // value every worker in the deployment holds. + Expect(plaintext).ToNot(Equal("shared-registration-token")) + + node, err := registry.Get(context.Background(), resp["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + // Stored as a hash, never as the secret. + Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext))) + Expect(node.TunnelTokenHash).ToNot(Equal(plaintext)) + // And it is a DIFFERENT column from the registration token's hash, + // which is what the tunnel used to compare against. + Expect(node.TunnelTokenHash).ToNot(Equal(node.TokenHash)) + Expect(node.TokenHash).To(Equal(hashOf("shared-registration-token"))) + + // The security property, which none of the above actually pins: a + // second node registering with the SAME shared token gets a + // DIFFERENT credential. Everything above is satisfied by a secret + // derived deterministically from the registration token, which + // would isolate nothing; a mutation that did exactly that passed + // every assertion before this one. + other := register(`{"name":"worker-tunnel-2","address":"10.0.0.3:50052","token":"shared-registration-token"}`, + "shared-registration-token", true) + Expect(other["tunnel_token"]).ToNot(Equal(plaintext)) + }) + + It("rotates the tunnel credential on every re-registration", func() { + body := `{"name":"worker-rotate","address":"10.0.0.4:50051"}` + first := register(body, "", true) + second := register(body, "", true) + + Expect(second["id"]).To(Equal(first["id"]), "re-registration must keep the node identity") + firstToken := first["tunnel_token"].(string) + secondToken := second["tunnel_token"].(string) + // Only the hash is stored, so a re-registering worker cannot be told + // the secret it already holds; the alternative to rotating would be + // storing the plaintext. + Expect(secondToken).ToNot(Equal(firstToken)) + + node, err := registry.Get(context.Background(), first["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(Equal(hashOf(secondToken))) + Expect(node.TunnelTokenHash).ToNot(Equal(hashOf(firstToken))) + }) + + It("issues a tunnel credential to a node still awaiting approval", func() { + resp := register(`{"name":"worker-pending","address":"10.0.0.5:50051"}`, "", false) + Expect(resp["status"]).To(Equal(nodes.StatusPending)) + // Deliberately unlike the agent API key and the NATS JWT, which are + // both withheld from a pending node. Those work the moment they are + // issued; this one does not, because the tunnel endpoint re-reads + // the node's status on every dial and refuses a pending node. A + // worker that registers exactly once would otherwise never receive + // one, since approval alone prompts no re-registration. + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + }) + It("returns nats_jwt when account seed is configured", func() { akp, err := nkeys.CreateAccount() Expect(err).ToNot(HaveOccurred()) diff --git a/core/services/cluster/tunnelproto.go b/core/services/cluster/tunnelproto.go new file mode 100644 index 000000000000..39f5014b50b4 --- /dev/null +++ b/core/services/cluster/tunnelproto.go @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + "strings" +) + +// The framing every stream on a worker tunnel opens with. +// +// A yamux stream on its own carries no destination: the frontend opens one and +// the worker has to be told what it is for. So the first thing on every stream +// is a request frame naming a TAG (which local service) and a TARGET (which +// instance of it), and the worker answers with a reply frame before either side +// speaks the tunnelled protocol. +// +// The reply is not optional and is not sent only on failure, which is the part +// that is easy to get wrong. The protocols carried here are client-speaks-first +// (gRPC sends an HTTP/2 preface, HTTP sends a request line), so a reply sent +// only when the worker refuses would arrive interleaved with a response body on +// the streams that succeeded, and the frontend would have no safe moment to +// look for it. Always sending one costs a round trip per stream, which is paid +// once per pooled connection rather than once per request. +// +// Both frames are length-prefixed rather than newline-delimited so a reader +// consumes exactly the header and not one byte of what follows: the stream is +// handed to gRPC or net/http afterwards, and a buffered reader that over-read +// would eat the beginning of their conversation. + +const ( + // StreamTagGRPC routes a stream to a backend process on the worker. Its + // target is the address that backend listens on, which the worker resolves + // itself; see the worker's tunnel services for what it will accept. + StreamTagGRPC = "grpc" + + // StreamTagHTTP routes a stream to the worker's own HTTP server, the one + // that serves file staging and backend logs. Its target is ignored: there + // is exactly one such server per worker and only the worker knows where it + // bound. + StreamTagHTTP = "http" +) + +// maxTunnelFrame bounds a header frame. It is a defence against a peer that +// declares a huge length and never sends it, not a size the protocol needs: +// the longest real frame is a tag plus a host:port, well under a hundred +// bytes. A reader that refuses early cannot be made to allocate on demand. +const maxTunnelFrame = 1024 + +// The reply codes. They travel on the wire, so they are strings rather than +// integers: a frontend reading a code from a worker it does not recognise can +// at least log something an operator can search for. +const ( + replyAccepted = "ok" + replyCodeUnknownTag = "unknown-tag" + replyCodeUnavailable = "unavailable" + replyCodeBadRequest = "bad-request" + replyPrefixRefused = "err " + streamRequestSeparator = " " +) + +// The three refusals a worker can send, kept apart on purpose. +// +// This is the phase's standing rule in its wire form. An unknown tag is a fact +// about what this worker SERVES and will not change until the worker is +// upgraded; an unavailable target is an infrastructure failure that may well +// succeed on the next attempt; a bad request is this frontend's own bug. A +// caller retries the second, gives up on the first, and reports the third. +// Collapsing them into one error would make a frontend retry a stream that can +// never work, or abandon a backend that was merely restarting. +// +// None of them wraps a node-absence error, and none must ever be built over +// one: a refusal is proof the worker is CONNECTED and answered. +var ( + ErrStreamTagUnknown = errors.New("cluster: the worker does not serve that stream tag") + ErrStreamTargetUnavailable = errors.New("cluster: the worker could not reach the local service for that stream") + ErrStreamRequestInvalid = errors.New("cluster: the worker rejected the stream request as malformed") +) + +// WriteStreamRequest sends the opening frame naming what the stream is for. +// +// An empty tag is refused here rather than on the wire, because the worker +// would answer it with ErrStreamRequestInvalid and the caller would learn a +// round trip later what it could have been told at once. +func WriteStreamRequest(w io.Writer, tag, target string) error { + if tag == "" { + return fmt.Errorf("writing a tunnel stream request: empty tag") + } + if strings.Contains(tag, streamRequestSeparator) { + // The separator is a single space and the split is on the FIRST one, so + // a tag containing a space would silently move part of itself 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 +// carries no argument. +// +// A malformed frame is returned as an ordinary error, NOT as +// ErrStreamRequestInvalid: that sentinel is what a worker SENDS to describe a +// refusal, and a reader that produced it here would leave a caller unable to +// tell "the peer refused my request" from "I could not read the peer's". +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 "", "", fmt.Errorf("reading a tunnel stream request: empty tag") + } + return tag, target, nil +} + +// WriteStreamAccepted tells the frontend the stream is now carrying the +// tunnelled protocol. Everything after this frame belongs to that protocol. +func WriteStreamAccepted(w io.Writer) error { + return writeFrame(w, replyAccepted) +} + +// WriteStreamRefusal reports why a stream will not be served. The caller closes +// the stream afterwards; this only says why. +// +// An unrecognised reason is sent as bad-request with its text attached rather +// than being dropped, because a refusal a frontend cannot read is +// indistinguishable from a worker that hung up, and those are different +// problems. +func WriteStreamRefusal(w io.Writer, reason error) error { + code := replyCodeBadRequest + switch { + case errors.Is(reason, ErrStreamTagUnknown): + code = replyCodeUnknownTag + case errors.Is(reason, ErrStreamTargetUnavailable): + code = replyCodeUnavailable + case errors.Is(reason, ErrStreamRequestInvalid): + code = replyCodeBadRequest + } + + text := "" + if reason != nil { + text = strings.Map(func(r rune) rune { + // The frame is length-prefixed so a newline would not corrupt it, + // but this text reaches a log line on the far side and a cause + // spanning lines is what makes one unsearchable. + if r == '\n' || r == '\r' { + return ' ' + } + return r + }, reason.Error()) + } + frame := replyPrefixRefused + code + streamRequestSeparator + text + if len(frame) > maxTunnelFrame { + frame = frame[:maxTunnelFrame] + } + return writeFrame(w, frame) +} + +// ReadStreamReply reads the worker's answer. nil means the stream is now +// carrying the tunnelled protocol. +// +// A failure to READ the reply is returned as itself, never as one of the +// refusal sentinels. The distinction is the point of this function: a refusal +// means the worker is connected and said no, while a read failure means the +// tunnel broke, and a caller that treated the second as the first would report +// a dead link as a policy decision. +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) + switch code { + case replyCodeUnknownTag: + return fmt.Errorf("%w: %s", ErrStreamTagUnknown, text) + case replyCodeUnavailable: + return fmt.Errorf("%w: %s", ErrStreamTargetUnavailable, text) + case replyCodeBadRequest: + return fmt.Errorf("%w: %s", ErrStreamRequestInvalid, text) + default: + // A code from a newer worker. Reported as an error carrying the code + // rather than mapped onto the nearest known one, so a frontend does not + // retry forever against a refusal that means something else entirely. + return fmt.Errorf("tunnel stream refused with unrecognised code %q: %s", code, text) + } +} + +// writeFrame writes one length-prefixed frame in a single Write. +// +// One Write, not two: the underlying stream is a yamux stream whose writes +// become discrete data frames, and splitting the length from the payload would +// put the reader one frame away from a header for no reason. It also keeps the +// adapter in wsconn.go to one WebSocket message per frame. +func writeFrame(w io.Writer, payload string) error { + if len(payload) > maxTunnelFrame { + return fmt.Errorf("tunnel frame is %d bytes, over the %d-byte limit", len(payload), maxTunnelFrame) + } + buf := make([]byte, 2+len(payload)) + binary.BigEndian.PutUint16(buf[:2], uint16(len(payload))) + copy(buf[2:], payload) + _, err := w.Write(buf) + return err +} + +// readFrame reads one length-prefixed frame. +// +// io.ReadFull rather than Read: a yamux stream returns whatever has arrived, +// and a header split across two data frames is ordinary rather than +// exceptional. It also converts a truncated frame into io.ErrUnexpectedEOF, +// which is what a peer that hung up mid-header should look 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) > maxTunnelFrame { + return "", fmt.Errorf("tunnel frame declares %d bytes, over the %d-byte limit", n, maxTunnelFrame) + } + 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/nodes/registry.go b/core/services/nodes/registry.go index e147a995574d..dd4ae23778a2 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -21,15 +21,27 @@ import ( // Workers are generic — they don't have a fixed backend type. // The SmartRouter dynamically installs backends via NATS backend.install events. type BackendNode struct { - 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 - 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 + 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 + // TunnelTokenHash is the SHA-256 of this node's OWN tunnel credential, the + // one it presents at GET /api/cluster/connect. It is not the registration + // token: registration mints a fresh random secret per node, returns the + // plaintext once, and stores only this hash, so a leaked registration token + // no longer opens a tunnel for every node whose ID an attacker can read. + // + // Empty means no tunnel credential has been minted for this node yet, which + // is what a node registered by an older LocalAI looks like. Such a node + // cannot tunnel until it registers again. That is deliberate: the column + // cannot be back-filled, because the plaintext 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 // 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 @@ -664,6 +676,24 @@ func (r *NodeRegistry) UpdateAuthRefs(ctx context.Context, nodeID, authUserID, a }).Error } +// SetTunnelTokenHash records the hash of a freshly minted tunnel credential for +// a node, replacing whatever was there. +// +// Replacing is the whole design and not an accident of the implementation. Only +// the hash is stored, so a re-registering worker cannot be told the secret it +// already has, and the alternative to rotating would be storing the plaintext. +// The live tunnel of a worker that re-registers is unaffected, because the +// credential is checked when a tunnel is DIALLED and never again; what changes +// is which secret its next reconnect must present, and the worker learns that +// in the same response that rotated it. +func (r *NodeRegistry) SetTunnelTokenHash(ctx context.Context, nodeID, hash string) error { + // Not Updates(struct): a struct update zero-skips, so this could never + // clear the column, and a caller that means to clear it would be silently + // ignored. + return r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", nodeID). + Update("tunnel_token_hash", hash).Error +} + // ApproveNode sets a pending node's status to healthy. func (r *NodeRegistry) ApproveNode(ctx context.Context, nodeID string) error { result := r.db.WithContext(ctx).Model(&BackendNode{}). diff --git a/core/services/worker/config.go b/core/services/worker/config.go index 8057e69fe790..6a3a19b39dbc 100644 --- a/core/services/worker/config.go +++ b/core/services/worker/config.go @@ -58,7 +58,12 @@ type Config struct { RegistrationRequireAuth bool `env:"LOCALAI_REGISTRATION_REQUIRE_AUTH" default:"false" help:"Refuse to start the HTTP file-transfer server when no registration token is set (otherwise it fails open and serves read/write to models/staging/data unauthenticated)" group:"registration"` DistributedRequireAuth bool `env:"LOCALAI_DISTRIBUTED_REQUIRE_AUTH" default:"false" help:"Umbrella switch implying both --nats-require-auth and --registration-require-auth" group:"distributed"` HeartbeatInterval string `env:"LOCALAI_HEARTBEAT_INTERVAL" default:"10s" help:"Interval between heartbeats" group:"registration"` - NodeLabels string `env:"LOCALAI_NODE_LABELS" help:"Comma-separated key=value labels for this node (e.g. tier=fast,gpu=a100)" group:"registration"` + // WorkerTunnel holds one outbound multiplexed connection to the frontend + // and serves the frontend's requests over it, so the worker needs no + // inbound port. Turning it off leaves the worker reachable only at the + // addresses it advertises, which is the pre-tunnel behaviour. + WorkerTunnel bool `env:"LOCALAI_WORKER_TUNNEL" default:"true" help:"Hold one outbound multiplexed tunnel to the frontend and serve its requests over it, so this worker needs no inbound port." group:"distributed"` + NodeLabels string `env:"LOCALAI_NODE_LABELS" help:"Comma-separated key=value labels for this node (e.g. tier=fast,gpu=a100)" group:"registration"` // MaxReplicasPerModel caps how many replicas of any one model can run on // this worker concurrently. Default 1 = historical single-replica // behavior. Set higher when a node has enough VRAM to host multiple diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go new file mode 100644 index 000000000000..5ab812d8c23e --- /dev/null +++ b/core/services/worker/tunnel.go @@ -0,0 +1,617 @@ +package worker + +import ( + "cmp" + "context" + "errors" + "fmt" + "math/rand/v2" + "net" + "net/http" + "net/url" + "strconv" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" + "github.com/libp2p/go-yamux/v5" + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +// The worker end of the tunnel. +// +// The worker DIALS OUT and never listens. It holds one WebSocket to the +// frontend load balancer, multiplexed with yamux, and every request the +// frontend makes of this worker arrives as a stream inside it. That is the +// whole point: a worker behind NAT, in another cluster or on a laptop needs no +// inbound port and no reachable address. +// +// This side is the yamux CLIENT and it only ACCEPTS streams; the frontend is +// the server and it only opens them. The frontend asks, the worker answers. +// Nothing here opens a stream, and a stream this side opened would park on the +// frontend's accept backlog, which accepts none. + +const ( + // tunnelBackoffBase is the shortest wait between reconnects, before jitter. + tunnelBackoffBase = 500 * time.Millisecond + + // tunnelBackoffMax is the ceiling on that wait. + // + // The ceiling is the interesting half. Without one, a worker that sits + // through a long frontend outage backs off into hours and does not come + // back for a long time after the frontend does; with one, the worst case + // for rejoining is bounded by this. The floor and the jitter are what stop + // a fleet of workers from turning a rolling restart into a retry storm + // against the first replica back up. + tunnelBackoffMax = 30 * time.Second + + // tunnelHealthyAfter is how long a session must last before the backoff is + // allowed back to its floor. + // + // Resetting on CONNECT rather than on a session that lasted is the classic + // way to build a reconnect storm that looks like a backoff: during a + // rolling restart a replica accepts the dial and dies moments later, so + // every attempt "succeeds" and every wait is the floor. This is set to the + // yamux keepalive interval, which is the shortest interval over which a + // session that is merely up can be told from one that is working. + tunnelHealthyAfter = 30 * time.Second + + // tunnelHandshakeTimeout bounds the WebSocket upgrade, matching the peer + // link's. + tunnelHandshakeTimeout = 10 * time.Second + + // tunnelHeaderTimeout bounds how long a stream may go without sending the + // request frame that says what it is for. It is generous because it bounds + // only the frontend's own framing, which it writes immediately after + // opening the stream; it is present because without it a stream that sends + // nothing holds a goroutine and one of the session's stream slots for as + // long as the tunnel lives. + tunnelHeaderTimeout = 15 * time.Second +) + +// LocalService opens a connection to one service running on this worker. +// +// target is the tag-specific argument from the stream's request frame, and the +// service decides what it will accept: the frontend naming an address does not +// oblige the worker to dial it. See loopbackService, which is what the worker +// actually installs. +type LocalService func(ctx context.Context, target string) (net.Conn, error) + +// TunnelConfig configures the tunnel a worker holds to the frontend. +type TunnelConfig struct { + // FrontendURL is the same value the worker registers against + // (LOCALAI_REGISTER_TO). Its scheme is mapped to ws/wss here. + FrontendURL string + + // NodeID is the identity registration assigned this worker. + NodeID string + + // Token supplies the node's own tunnel credential. + // + // A function and not a string, and that is load-bearing rather than + // stylistic. The credential is re-minted on every registration, so a client + // that captured one at startup would keep presenting a value the frontend + // stopped accepting the moment anything re-registered this worker, and + // would lock itself out with no way back. It is called once per DIAL. + Token func() string + + // Services routes an accepted stream by the tag in its request frame. A tag + // with no entry here is refused; see Tunnel.accept. + Services map[string]LocalService + + // Seams the specs replace. They are unexported so they are not part of the + // package's API: a caller cannot reach them, and the internal test file can. + sleep func(ctx context.Context, d time.Duration) error + now func() time.Time + headerTimeout time.Duration +} + +// Tunnel is a running worker tunnel: one goroutine holding one session at a +// time, reconnecting when it dies, until Close. +type Tunnel struct { + endpoint string + nodeID string + token func() string + services map[string]LocalService + dialer *websocket.Dialer + + headerTimeout time.Duration + sleep func(ctx context.Context, d time.Duration) error + // now measures how long a session lasted, and nothing else. Deadlines are + // taken from time.Now directly: a spec that fakes this clock to exercise + // the backoff must not thereby move every I/O deadline in the package. + now func() time.Time + + cancel context.CancelFunc + done chan struct{} + closeOnce sync.Once +} + +// StartTunnel dials the frontend and holds the tunnel until ctx is cancelled or +// Close is called. +// +// The returned error is about this CONFIGURATION, never about the frontend. A +// frontend that is down, that has not been upgraded, or that refuses the +// credential is not a reason for a worker to fail to start: it retries, with +// backoff, in the background. Failing to start on a dial would make a frontend +// restart into a fleet-wide worker outage. +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 the tunnel's routing table cannot change under the accept loop + // after it has started. + services := make(map[string]LocalService, len(cfg.Services)) + for tag, svc := range cfg.Services { + services[tag] = svc + } + + t := &Tunnel{ + endpoint: endpoint, + nodeID: cfg.NodeID, + token: cfg.Token, + services: services, + dialer: &websocket.Dialer{ + HandshakeTimeout: tunnelHandshakeTimeout, + // A worker reaches its frontend over the public internet in the + // deployments this exists for, so unlike the replica-to-replica + // peer link this DOES honour the environment's proxy settings. + Proxy: http.ProxyFromEnvironment, + }, + headerTimeout: cmp.Or(cfg.headerTimeout, tunnelHeaderTimeout), + 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 finish. It is idempotent. +func (t *Tunnel) Close() error { + t.closeOnce.Do(func() { + t.cancel() + <-t.done + }) + return nil +} + +// run holds one session at a time, reconnecting with bounded backoff. +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 evidence the frontend is healthy. + // See tunnelHealthyAfter for why "we connected" 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, serves streams until the session ends, and leaves +// nothing running behind it. +func (t *Tunnel) connectAndServe(ctx context.Context) error { + ws, err := t.dial(ctx) + if err != nil { + return err + } + + sess, err := yamux.Client(cluster.WebsocketConn(ws), nil, nil) + if err != nil { + _ = ws.Close() + return fmt.Errorf("starting the worker tunnel session: %w", err) + } + xlog.Info("Worker tunnel established", "node", t.nodeID, "frontend", t.endpoint) + + // Streams are served under a context of the SESSION's, not the loop's. A + // stream goroutine parked in a local dial would otherwise outlive the + // session it belongs to and hold the reconnect below behind it. + sessCtx, endSession := context.WithCancel(ctx) + + // AcceptStream takes no context, so something else has to break it when the + // worker is shutting down; closing the session is that something. + watchdogDone := make(chan struct{}) + go func() { + defer close(watchdogDone) + select { + case <-sessCtx.Done(): + _ = sess.Close() + case <-sess.CloseChan(): + } + }() + + var streams sync.WaitGroup + serveErr := t.serve(sessCtx, sess, &streams) + + endSession() + _ = sess.Close() + <-watchdogDone + // Closing the session unblocks every stream goroutine: Session.close walks + // its stream table calling forceClose on each (session.go:334-338), and + // forceClose puts both directions in halfReset and calls notifyWaiting + // (stream.go:371-388), which wakes a parked Read and fails a parked Write. + // Waiting here is what keeps a reconnect from overlapping the streams of + // the session it replaced. + streams.Wait() + return serveErr +} + +// serve accepts streams until the session ends. +// +// One goroutine per stream, and an error from a stream never reaches this loop. +// A single malformed or unroutable request must not cost this worker every +// other request in flight on the same session. +func (t *Tunnel) serve(ctx context.Context, sess *yamux.Session, streams *sync.WaitGroup) error { + for { + stream, err := sess.AcceptStream() + if err != nil { + return err + } + streams.Add(1) + go func() { + defer streams.Done() + t.handleStream(ctx, stream) + }() + } +} + +// handleStream reads one stream's request frame and either splices it 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, and here that is + // not a figure of speech: nothing supervises this goroutine, so an + // unrecovered panic ends the PROCESS, which ends the session and every + // other stream on it. Unlike the frontend's handler next door this does not + // re-panic, because there is no recovery middleware above it to report the + // panic; re-panicking would only be the crash. + // + // It covers what runs ON THIS goroutine: reading the request frame, the + // route lookup, and the local service's dial, which is the one of the three + // that runs caller-supplied code. It does NOT cover a panic inside Splice's + // own copy goroutines, which no recover here can reach. + 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 already answered and closed the stream. + return + } + + // Splice owns closing both ends from here. + if err := cluster.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 +// second result is false when the stream was refused, in which case the refusal +// has been sent and the stream closed. +func (t *Tunnel) accept(ctx context.Context, stream net.Conn) (net.Conn, bool) { + // Deliberately time.Now and not t.now: this is an I/O deadline, and the + // clock seam exists only to measure how long a session lasted. + if err := stream.SetReadDeadline(time.Now().Add(t.headerTimeout)); err != nil { + // Nothing is readable on a stream whose deadline cannot be set, so this + // is reported as an infrastructure failure rather than pushed past. + t.refuse(stream, fmt.Errorf("%w: arming the request deadline: %v", cluster.ErrStreamTargetUnavailable, err)) + return nil, false + } + + tag, target, err := cluster.ReadStreamRequest(stream) + if err != nil { + // Includes the deadline above expiring. Both are "this stream never + // told me what it wanted", which is the frontend's problem to fix, not + // something a retry against this worker resolves. + t.refuse(stream, fmt.Errorf("%w: %v", cluster.ErrStreamRequestInvalid, err)) + return nil, false + } + + svc, known := t.services[tag] + if !known { + // A ROUTING fact about this worker, and it is reported as itself. A + // frontend that reads this knows a retry is pointless until the worker + // is upgraded, which is not what it should conclude from the + // unavailable below. + t.refuse(stream, fmt.Errorf("%w: %q", cluster.ErrStreamTagUnknown, tag)) + return nil, false + } + + // Cleared before the local dial rather than after the reply: everything + // past the request frame belongs to the tunnelled protocol, which brings + // its own deadlines, and one left armed here would abort a long inference + // stream in the middle. + if err := stream.SetReadDeadline(time.Time{}); err != nil { + t.refuse(stream, fmt.Errorf("%w: clearing the request deadline: %v", cluster.ErrStreamTargetUnavailable, err)) + return nil, false + } + + local, err := svc(ctx, target) + if err != nil { + // An INFRASTRUCTURE failure, which a frontend may retry, and which must + // never be reported as the unknown tag above. + t.refuse(stream, fmt.Errorf("%w: %v", cluster.ErrStreamTargetUnavailable, err)) + return nil, false + } + + if err := cluster.WriteStreamAccepted(stream); err != nil { + // The frontend never learns the stream was accepted, so it cannot be + // used; closing the local connection here is what stops an accepted + // backend connection leaking per failed reply. + xlog.Debug("worker tunnel could not accept a stream", "node", t.nodeID, "error", err) + _ = local.Close() + _ = stream.Close() + return nil, false + } + return local, true +} + +// refuse reports why a stream will not be served and then ENDS it. +// +// The close is the part that matters and it is not optional. A worker that says +// why and leaves the stream open has parked the frontend on a request that will +// never be answered, which reads as a slow worker rather than a refused +// request, and a deadline on the far side cannot tell those apart. The reply is +// what makes the refusal legible; the close is what makes it prompt. +// +// The reply is therefore best-effort and the close is not: a reply that could +// not be written still gets the stream closed. +func (t *Tunnel) refuse(stream net.Conn, reason error) { + if err := cluster.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 and returns it. +func (t *Tunnel) dial(ctx context.Context) (*websocket.Conn, error) { + // Read HERE, once per dial. See TunnelConfig.Token. + token := t.token() + if token == "" { + // Not a dial that fails with "unauthorized": this worker has no + // credential yet, which is a different condition from the frontend + // rejecting one, and an operator reading "unauthorized" would go + // looking for a token 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) + + ws, resp, err := t.dialer.DialContext(ctx, t.endpoint, header) + if err != nil { + if resp != nil { + // gorilla reports every non-101 as the same ErrBadHandshake, so + // 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) + } + return ws, nil +} + +// tunnelDialError carries the HTTP status a refused dial came back with, so the +// four refusals the frontend can give are not logged as one. +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 } + +// logSessionEnded says why the tunnel is reconnecting, at a level that matches +// what the operator can do about it. +// +// The distinctions are the point rather than decoration. "Awaiting approval" +// and "your token is wrong" and "this frontend does not do tunnels" send an +// operator to three different places, and a worker retries all three the same +// way: none of them is a reason to stop, because a re-registration or an admin +// action fixes each without restarting the worker. +func (t *Tunnel) logSessionEnded(err error, attempt int, delay time.Duration) { + if err == nil { + xlog.Info("Worker tunnel closed, reconnecting", "node", t.nodeID, "attempt", attempt, "retry_in", delay) + return + } + + var dialErr *tunnelDialError + if errors.As(err, &dialErr) { + switch dialErr.status { + case http.StatusUnauthorized: + xlog.Warn("Frontend rejected this worker's tunnel credential; re-registering will mint a fresh one", + "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, reconnecting", "node", t.nodeID, "attempt", attempt, "retry_in", delay, "error", err) +} + +// tunnelBackoffDelay returns how long to wait before reconnect attempt n. +// +// Equal jitter: half the delay is fixed and half is drawn. Full jitter, which +// draws over the whole interval, can produce a near-zero wait, and a worker +// that can draw a near-zero wait can spin; keeping a floor means no single +// worker ever does, while the drawn half is what stops a fleet that all lost +// the same replica from resynchronising onto the same instant. +func tunnelBackoffDelay(attempt int) time.Duration { + if attempt < 1 { + attempt = 1 + } + d := tunnelBackoffMax + // The shift is guarded twice. The bound on attempt keeps the shift itself + // defined, and the positivity check catches the overflow that would + // otherwise turn a long outage into a NEGATIVE delay, which is a tight loop + // wearing a backoff's costume. + if attempt <= 40 { + if scaled := tunnelBackoffBase << (attempt - 1); scaled > 0 && scaled < tunnelBackoffMax { + d = scaled + } + } + return d/2 + time.Duration(rand.Int64N(int64(d/2)+1)) +} + +// tunnelSleep waits for d, or returns early when ctx is cancelled. +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 frontend URL a worker registers against into the +// WebSocket URL it dials its tunnel on. +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 rather than assigned, so a frontend served under a path prefix + // keeps it. Registration builds its URLs the same way. + u.Path = strings.TrimRight(u.Path, "/") + cluster.ConnectPath + u.RawQuery = url.Values{"id": []string{nodeID}}.Encode() + return u.String(), nil +} + +// loopbackService routes a tagged stream to a process listening on this +// worker's own loopback interface. +// +// The HOST the frontend names is discarded and only the port is used, which is +// deliberate and is the security property this function exists for. A tunnel +// terminates inside the worker process, so a stream arriving on it can reach +// anything the worker can reach; without this, whoever holds the frontend end +// could make every worker in the fleet dial arbitrary hosts on its private +// network, turning the tunnel into a proxy into the worker's LAN. Discarding +// the host reduces the reachable set to this machine. +// +// It also happens to be what makes the tunnel work BEFORE the workers stop +// advertising themselves: today the frontend names the address the worker +// registered, which is a routable one, and after that change it will name a +// loopback one. Both resolve to the same place here. +// +// The port range is the one the worker's own port allocator hands to backend +// processes, so a stream cannot be pointed at some unrelated service that +// happens to be listening on this host. +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("routing a tunnel stream: %q is not a host:port: %w", target, err) + } + port, err := strconv.Atoi(portStr) + if err != nil { + return nil, fmt.Errorf("routing a tunnel stream: %q has no numeric port: %w", target, err) + } + if port < minPort || port > maxPort { + return nil, fmt.Errorf("routing a tunnel stream: port %d is outside this worker's backend range [%d, %d]", port, minPort, maxPort) + } + var d net.Dialer + return d.DialContext(ctx, "tcp", net.JoinHostPort("127.0.0.1", portStr)) + } +} + +// fixedService routes a tagged stream to one address on this worker, ignoring +// whatever the frontend named. +// +// There is exactly one HTTP server per worker and only the worker knows where +// it bound, so the frontend has nothing useful to say about the target and is +// not given the chance to say it. +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 rewrites a bind address into one that reaches the same listener +// from inside this process. +// +// A server bound to 0.0.0.0 is reachable on loopback, but dialling 0.0.0.0 is +// only accidentally equivalent to dialling localhost and is not on every +// platform, so the wildcard is replaced rather than 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("127.0.0.1", port) + } + return bindAddr +} diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go new file mode 100644 index 000000000000..09285c401591 --- /dev/null +++ b/core/services/worker/tunnel_test.go @@ -0,0 +1,509 @@ +package worker + +import ( + "context" + "encoding/binary" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +// awaitErr runs fn on its own goroutine and reports its result on a channel. +// +// Every blocking read in this file goes through it, and that is the single most +// load-bearing decision in the whole suite. The obvious way to assert "the +// worker refused this stream promptly" is to arm a read deadline and expect an +// error, and phase 1 shipped exactly that in three places: it held in none, +// because a stream the worker never answers AT ALL satisfies a deadline +// assertion just as well as one it refused. Reading with NO deadline, on +// another goroutine, and asserting the channel delivers, inverts that: a parked +// stream delivers nothing and the Eventually 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 +} + +// fakeFrontend is the far side of the tunnel: it speaks the real WebSocket +// upgrade and the real yamux server handshake, so these specs exercise the +// wire, not a mock of it. It is deliberately NOT core/http's handler; that one +// needs a database, and what is under test here is the client. +type fakeFrontend struct { + srv *httptest.Server + sessions chan *yamux.Session + dials chan tunnelDial + + // closeAtOnce makes every accepted session die immediately, which is what a + // frontend replica going down during a rolling restart looks like from the + // worker. + closeAtOnce bool +} + +func newFakeFrontend(closeAtOnce bool) *fakeFrontend { + f := &fakeFrontend{ + sessions: make(chan *yamux.Session, 64), + dials: make(chan tunnelDial, 256), + closeAtOnce: closeAtOnce, + } + upgrader := websocket.Upgrader{} + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != cluster.ConnectPath { + w.WriteHeader(http.StatusNotFound) + return + } + select { + case f.dials <- tunnelDial{ + token: strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer "), + nodeID: r.URL.Query().Get("id"), + }: + default: + } + + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := yamux.Server(cluster.WebsocketConn(ws), nil, nil) + if err != nil { + _ = ws.Close() + return + } + if f.closeAtOnce { + _ = sess.Close() + return + } + select { + case f.sessions <- sess: + default: + _ = sess.Close() + } + })) + return f +} + +func (f *fakeFrontend) close() { + for { + select { + case sess := <-f.sessions: + _ = sess.Close() + default: + f.srv.Close() + return + } + } +} + +// echoListener is a stand-in for a backend gRPC process on the worker: a local +// TCP listener that reads and writes back. +func echoListener() net.Listener { + ln, err := net.Listen("tcp", "127.0.0.1:0") + 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 +} + +// dialLocalTCP is the simplest possible LocalService: connect to whatever 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 + tunnel *Tunnel + ) + + BeforeEach(func() { + ctx, cancel = context.WithCancel(context.Background()) + }) + + AfterEach(func() { + if tunnel != nil { + Expect(tunnel.Close()).To(Succeed()) + tunnel = nil + } + cancel() + if frontend != nil { + frontend.close() + frontend = nil + } + }) + + // start brings up the client against the fake frontend already created. + start := func(mutate func(*TunnelConfig)) { + 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 + tunnel, err = StartTunnel(ctx, cfg) + Expect(err).ToNot(HaveOccurred()) + } + + // session waits for the frontend to have accepted the worker's dial. + session := func() *yamux.Session { + var sess *yamux.Session + EventuallyWithOffset(1, frontend.sessions, "10s").Should(Receive(&sess)) + return sess + } + + 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[cluster.StreamTagGRPC] = dialLocalTCP + }) + + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(stream, cluster.StreamTagGRPC, ln.Addr().String())).To(Succeed()) + + reply := awaitErr(func() error { return cluster.ReadStreamReply(stream) }) + Eventually(reply, "10s").Should(Receive(BeNil())) + + _, err = stream.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + + buf := make([]byte, 4) + read := awaitErr(func() error { + _, err := io.ReadFull(stream, buf) + return err + }) + Eventually(read, "10s").Should(Receive(BeNil())) + Expect(string(buf)).To(Equal("ping")) + }) + }) + + Describe("refusing a stream it cannot serve", func() { + // The refusal specs all read with NO deadline, on another goroutine. + // See awaitErr: a deadline would be satisfied by a stream that was + // merely parked, which is the exact defect this phase inherited. + + It("refuses an unknown tag promptly, and the stream ENDS rather than parking", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[cluster.StreamTagGRPC] = dialLocalTCP + }) + + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(stream, "no-such-tag", "")).To(Succeed()) + + // Two facts, in order, on one goroutine: the worker SAID why, and + // then the stream ended. A worker that only says why and leaves the + // stream open never sends on this channel, so the Eventually below + // fails rather than passing on a deadline. + type outcome struct{ reply, end error } + done := make(chan outcome, 1) + go func() { + var got outcome + got.reply = cluster.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(cluster.ErrStreamTagUnknown)) + Expect(got.end).To(MatchError(io.EOF)) + }) + + It("reports a local service it could not reach as unavailable, not as an unknown tag", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[cluster.StreamTagGRPC] = func(context.Context, string) (net.Conn, error) { + return nil, fmt.Errorf("connection refused") + } + }) + + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(stream, cluster.StreamTagGRPC, "127.0.0.1:1")).To(Succeed()) + + reply := awaitErr(func() error { return cluster.ReadStreamReply(stream) }) + var got error + Eventually(reply, "10s").Should(Receive(&got)) + // Distinct conditions must not be reported as each other: a caller + // gives up on an unknown tag and retries an unavailable target. + Expect(got).To(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(got).ToNot(MatchError(cluster.ErrStreamTagUnknown)) + }) + + It("ends a stream whose request never arrives instead of holding it open", func() { + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[cluster.StreamTagGRPC] = dialLocalTCP + c.headerTimeout = 50 * time.Millisecond + }) + + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + // Nothing is written. A worker that waits forever for a request it + // will never get holds a goroutine and a stream slot per dial. + ended := awaitErr(func() error { + _, err := io.Copy(io.Discard, stream) + return err + }) + Eventually(ended, "10s").Should(Receive(BeNil())) + }) + }) + + Describe("surviving a bad stream", func() { + It("keeps serving the session after one stream it could not read", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[cluster.StreamTagGRPC] = dialLocalTCP + }) + sess := session() + + // A frame that declares far more than it sends, then hangs up. The + // worker cannot parse it and must not take the session down with it. + 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.CloseWrite()).To(Succeed()) + badEnded := awaitErr(func() error { + _, err := io.Copy(io.Discard, bad) + return err + }) + Eventually(badEnded, "10s").Should(Receive(BeNil())) + + // Same session, a stream the worker can serve. + good, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(good, cluster.StreamTagGRPC, ln.Addr().String())).To(Succeed()) + reply := awaitErr(func() error { return cluster.ReadStreamReply(good) }) + Eventually(reply, "10s").Should(Receive(BeNil())) + + _, err = good.Write([]byte("still here")) + Expect(err).ToNot(HaveOccurred()) + buf := make([]byte, len("still here")) + read := awaitErr(func() error { + _, err := io.ReadFull(good, buf) + return err + }) + Eventually(read, "10s").Should(Receive(BeNil())) + Expect(string(buf)).To(Equal("still here")) + }) + It("serves streams concurrently, so one live stream does not block the next", func() { + ln := echoListener() + DeferCleanup(func() { _ = ln.Close() }) + + frontend = newFakeFrontend(false) + start(func(c *TunnelConfig) { + c.Services[cluster.StreamTagGRPC] = dialLocalTCP + }) + sess := session() + + // The first stream is accepted and then left open with nothing + // flowing, which is what an idle inference stream or a paused file + // transfer looks like. Serving streams from the accept loop rather + // than a goroutine each would park every later request behind it. + first, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(first, cluster.StreamTagGRPC, ln.Addr().String())).To(Succeed()) + firstReply := awaitErr(func() error { return cluster.ReadStreamReply(first) }) + Eventually(firstReply, "10s").Should(Receive(BeNil())) + + second, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(second, cluster.StreamTagGRPC, ln.Addr().String())).To(Succeed()) + secondReply := awaitErr(func() error { return cluster.ReadStreamReply(second) }) + Eventually(secondReply, "10s").Should(Receive(BeNil())) + }) + + 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[cluster.StreamTagGRPC] = dialLocalTCP + }) + sess := session() + + // Nothing supervises a stream goroutine, so an unrecovered panic + // here ends the process, which is the loudest possible way to kill + // the session. + boom, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(boom, "explodes", "")).To(Succeed()) + boomEnded := awaitErr(func() error { + _, err := io.Copy(io.Discard, boom) + return err + }) + Eventually(boomEnded, "10s").Should(Receive(BeNil())) + + good, err := sess.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(good, cluster.StreamTagGRPC, ln.Addr().String())).To(Succeed()) + reply := awaitErr(func() error { return cluster.ReadStreamReply(good) }) + Eventually(reply, "10s").Should(Receive(BeNil())) + }) + }) + + Describe("reconnecting", func() { + It("backs off exponentially between reconnects, bounded and never tight", func() { + frontend = newFakeFrontend(true) // every session dies at once + + 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: + } + return ctx.Err() + } + }) + + observed := make([]time.Duration, 0, 10) + for i := 0; i < 10; i++ { + var d time.Duration + Eventually(delays, "20s").Should(Receive(&d), "expected reconnect attempt %d", i+1) + observed = append(observed, d) + } + + for i, d := range observed { + // Never a tight loop: a worker that reconnect-storms a frontend + // during a rolling restart is a denial of service against the + // control plane. + Expect(d).To(BeNumerically(">", 0), "delay %d was not positive", i+1) + // Bounded: without a ceiling a worker that misses a rolling + // restart backs off into hours and never comes back. + Expect(d).To(BeNumerically("<=", tunnelBackoffMax), "delay %d exceeded the ceiling", i+1) + } + // And it actually grows. The jitter has a floor of half the + // unjittered delay, so the fourth attempt is at least 4x the base + // however the dice fall. + Expect(observed[3]).To(BeNumerically(">=", 4*tunnelBackoffBase)) + }) + + It("keeps backing off after a session that died at once", func() { + frontend = newFakeFrontend(true) + + 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: + } + return ctx.Err() + } + }) + + var last time.Duration + for i := 0; i < 6; i++ { + Eventually(delays, "20s").Should(Receive(&last)) + } + // A session that came up and died immediately is not evidence the + // frontend is healthy, so the delay must NOT be back at the floor. + Expect(last).To(BeNumerically(">", tunnelBackoffBase)) + }) + + It("returns to its shortest delay after a session that lasted", func() { + frontend = newFakeFrontend(true) + + // A clock that jumps a minute on every reading. The loop reads it + // once when a session comes up and once when it ends, so every + // session looks like it lasted a minute, which is longer than the + // threshold below which a session is not counted as healthy. + var ticks atomic.Int64 + base := time.Now() + + delays := make(chan time.Duration, 64) + start(func(c *TunnelConfig) { + c.now = func() time.Time { + return base.Add(time.Duration(ticks.Add(1)) * time.Minute) + } + c.sleep = func(ctx context.Context, d time.Duration) error { + select { + case delays <- d: + default: + } + return ctx.Err() + } + }) + + for i := 0; i < 6; i++ { + var d time.Duration + Eventually(delays, "20s").Should(Receive(&d), "expected reconnect attempt %d", i+1) + Expect(d).To(BeNumerically("<=", tunnelBackoffBase), + "delay %d did not return to the floor after a session that lasted", i+1) + } + }) + + It("presents the credential current at DIAL time, not the one it started with", 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() } + }) + + // Nothing survives a reconnect: the new owner replica has no record + // of the old session, and the worker's own credential may have been + // rotated by a re-registration in between. A client that captured + // its token once locks itself out on the first rotation. + 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")) + Expect(second.nodeID).To(Equal("node-1")) + }) + }) +}) diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 6434c3cd6b69..829472ffb7e8 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -16,6 +16,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/messaging" "github.com/mudler/LocalAI/core/services/nodes" grpc "github.com/mudler/LocalAI/pkg/grpc" @@ -94,13 +95,25 @@ func Run(ctx *cliContext.Context, cfg *Config) error { var ( nodeID string connectNats func() (*messaging.Client, error) + // tunnelToken reads the node's CURRENT tunnel credential. It is a + // function because the frontend rotates the credential on every + // registration, so the value a reconnect must present is not + // necessarily the one this worker started with. + tunnelToken func() string ) if cfg.NatsJWT != "" || cfg.NatsUserSeed != "" { - nid, _, _, _, regErr := regClient.RegisterWithRetry(shutdownCtx, registrationBody, 10) + res, regErr := regClient.RegisterFullWithRetry(shutdownCtx, registrationBody, 10) if regErr != nil { return fmt.Errorf("failed to register with frontend: %w", regErr) } - nodeID = nid + nodeID = res.ID + // This path registers exactly once and never again, so the credential + // it holds cannot go stale by rotation from its own side. It can still + // be superseded from OUTSIDE, by a second worker registering under the + // same node name, and there is nothing this worker can do about that + // but log the 401 and keep retrying. + staticTunnelToken := res.TunnelToken + tunnelToken = func() string { return staticTunnelToken } connectNats = func() (*messaging.Client, error) { return connectNATS(cfg.NatsURL, cfg.NatsJWT, cfg.NatsUserSeed, "", "", cfg.NatsAuthRequired(), natsTLS) } @@ -116,6 +129,10 @@ func Run(ctx *cliContext.Context, cfg *Config) error { return fmt.Errorf("failed to register with frontend: %w", regErr) } nodeID = res.ID + // The manager re-registers to refresh NATS credentials, and every + // registration rotates the tunnel credential too, so this reads the + // manager rather than capturing a value. + tunnelToken = credMgr.TunnelToken connectNats = func() (*messaging.Client, error) { var opts []messaging.Option if credMgr.HasCredentials() { @@ -163,6 +180,42 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // used to remove them, so a long-lived worker filled its own disk. StartEphemeralStagingCleanup(shutdownCtx, stagingDir, 0, 0) + // The tunnel is started here, after the HTTP server it fronts is listening + // and before any backend process exists. Both orders are deliberate: a + // stream tagged for HTTP that arrived before the server bound would be + // refused as unavailable, while a stream tagged for gRPC resolves its + // backend at dial time, so nothing has to exist yet for the tunnel to be + // useful. + // + // A failure to START it is fatal, unlike a failure to CONNECT: it means the + // frontend URL or this node's identity is unusable, and a worker that + // silently ran without its tunnel would look healthy while being + // unreachable to everything that dials through it. + tunnelBasePort := cfg.effectiveBasePort() + if cfg.WorkerTunnel { + tunnel, terr := StartTunnel(shutdownCtx, TunnelConfig{ + FrontendURL: cfg.RegisterTo, + NodeID: nodeID, + Token: tunnelToken, + Services: map[string]LocalService{ + // The frontend names a backend process by its port; the worker + // decides that only its own loopback, and only within its own + // backend port range, is reachable through it. + cluster.StreamTagGRPC: loopbackService(tunnelBasePort, cfg.effectiveMaxPort(tunnelBasePort)), + cluster.StreamTagHTTP: fixedService(loopbackAddr(httpAddr)), + }, + }) + if terr != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("starting the worker tunnel: %w", terr) + } + defer func() { + if err := tunnel.Close(); err != nil { + xlog.Warn("Closing the worker tunnel failed", "error", err) + } + }() + } + // Connect to NATS xlog.Info("Connecting to NATS", "url", sanitize.URL(cfg.NatsURL)) natsClient, err := connectNats() @@ -199,7 +252,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { }() // Process supervisor — manages multiple backend gRPC processes on different ports - basePort := cfg.effectiveBasePort() + basePort := tunnelBasePort // Buffered so NATS stop handler can send without blocking sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 1ad6822afd23..e276acf406cf 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -104,20 +104,43 @@ The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_R A worker can open one long-lived, multiplexed tunnel to the frontend instead of listening on a port of its own. It dials `GET /api/cluster/connect?id=`, the connection is upgraded to a WebSocket, and every subsequent request the frontend makes to that worker travels as a stream inside it. Nothing dials *into* the worker, so a worker behind NAT, in another Kubernetes cluster or on a laptop needs no inbound port and no reachable address. -The dial is authenticated against **the hash stored on that node's row, which today is still the registration token's hash**: the frontend hashes the presented bearer token and compares it with what registration recorded. A worker that presents a token belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A node still awaiting admin approval is refused with `403`. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity. +#### Each worker has its own tunnel credential -So the isolation is not there yet, and the caveat is worth stating plainly: because a worker registers by presenting the deployment's registration token, a leaked registration token plus a known node ID still gets a tunnel. What the tunnel endpoint does not do is trust the configured token directly, so when workers are issued their own per-node secrets and those secrets are what registration stores, the isolation becomes real with no change to this route. +The dial is authenticated against **that node's own tunnel credential**, which is not the registration token. Registration mints a fresh random secret per node, returns the plaintext once in the registration response as `tunnel_token`, and stores only its SHA-256. So a leaked registration token no longer opens a tunnel: an attacker who has it, and who knows a node ID, still cannot authenticate as that worker. -**Worker tunnels need `LOCALAI_REGISTRATION_TOKEN` set.** A worker registering against a frontend with no registration token configured sends no token, so nothing is stored on its row, and the tunnel has nothing to authenticate it against: every dial is refused with `401`. LocalAI warns about this at startup. Setting the token later is not enough on its own; the workers have to register again for the hash to be written. +A worker that presents a credential belonging to no node, or names a node ID the frontend has never seen, is refused with `401` before the WebSocket upgrade happens. A node still awaiting admin approval is refused with `403`. A frontend that cannot read its node table answers `500` rather than `401`, so a worker retries instead of re-registering under a new identity. + +**The credential is rotated on every registration.** That follows from storing only the hash: a re-registering worker cannot be told the secret it already holds, so it is given a new one. The worker's live tunnel is unaffected, because the credential is checked when a tunnel is *dialled* and never again; what changes is which secret the next reconnect presents, and the worker learns it in the same response that rotated it. + +**A node that has not registered since upgrading cannot tunnel.** Its row has no tunnel credential and the column cannot be back-filled, because the plaintext only ever existed in the response that minted it. Such a node is refused with `401` until it registers again, which a worker restart does. The frontend does *not* fall back to the registration token for these nodes. + +Unlike the agent worker's API key and its NATS credential, a tunnel credential **is** issued to a node still awaiting approval. It is inert until then: the tunnel route re-reads the node's status on every dial and refuses a pending one. Withholding it would instead strand workers that register exactly once, since approval on its own prompts no re-registration. + +A tunnel credential does not replace `LOCALAI_REGISTRATION_TOKEN`. Without one, node registration itself is unauthenticated, so anyone who can reach the frontend can register a worker and be issued a working tunnel credential for it. LocalAI warns about that at startup. The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims the tunnels it still holds on a live session as soon as it re-registers, skipping any whose socket has already gone. That re-claim needs the replica to have an advertised address: without one it never had an instance row to begin with, and its tunnels are usable only by the replica holding them. | Method | Path | Description | |--------|------|-------------| -| `GET` | `/api/cluster/connect?id=` | Worker opens its multiplexed tunnel (`Authorization: Bearer `) | +| `GET` | `/api/cluster/connect?id=` | Worker opens its multiplexed tunnel (`Authorization: Bearer `) | The route is exempt from the normal session/API-key authentication (it authenticates itself, like `/api/cluster/peer`) and is registered in every deployment. Outside distributed mode there is no node table to check a token against, so it answers `503`. +#### What the worker does with the tunnel + +The worker holds the tunnel with one goroutine: it dials, serves the frontend's streams until the session dies, and dials again. Every stream opens with a small frame naming which local service it is for, and the worker answers before either side speaks the tunnelled protocol: + +| Tag | Goes to | Target | +|-----|---------|--------| +| `grpc` | a backend process on this worker | the port; the host is discarded and only `127.0.0.1` is dialled, within the worker's own backend port range | +| `http` | the worker's own file-transfer and backend-log server | ignored; there is one such server and only the worker knows where it bound | + +A stream naming a tag the worker does not serve, or a local service it could not reach, is refused with a reason and the stream is **ended** rather than left open. Those two refusals are distinct on the wire on purpose: a frontend gives up on the first and may retry the second. One bad stream never affects the others or the session. + +Reconnects use exponential backoff with jitter: the interval doubles from 500ms up to a ceiling of 30 seconds, and each wait is drawn between half of that interval and all of it, so no worker ever spins and a fleet that lost the same replica does not come back in lockstep. The interval returns to its floor only after a session that lasted at least 30 seconds. That last part is what stops a rolling frontend restart, where every dial succeeds and then dies moments later, from turning a fleet of workers into a retry storm against the first replica back up. A worker that is refused (`401`, `403`) keeps retrying on the same schedule rather than exiting: a re-registration or an admin approval fixes both without restarting it. + +Set `LOCALAI_WORKER_TUNNEL=false` on a worker to turn the tunnel off and go back to the frontend dialling the worker's advertised addresses. + ### The model load deadline scales with the checkpoint The `LoadModel` deadline starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint read and pipeline init. That work is proportional to the bytes on disk, which makes any fixed deadline a model-size cliff rather than a timeout: a 70 GB video checkpoint on a Jetson Thor worker failed reproducibly against the old fixed 5m default (`rpc error: code = DeadlineExceeded` after 953.5s of wall clock, roughly 11m of which was backend install and staging), and simply raising the constant would only move the cliff to the next larger model while making a genuinely wedged *small* model hang for the whole inflated duration. @@ -341,6 +364,7 @@ local-ai worker \ | `--registration-require-auth` | `LOCALAI_REGISTRATION_REQUIRE_AUTH` | `false` | Refuse to start the HTTP file-transfer server when no registration token is set (it would otherwise fail open) | | `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | Umbrella switch implying both `--registration-require-auth` and `--nats-require-auth` | | `--heartbeat-interval` | `LOCALAI_HEARTBEAT_INTERVAL` | `10s` | Interval between heartbeat pings | +| `--worker-tunnel` | `LOCALAI_WORKER_TUNNEL` | `true` | Hold one outbound multiplexed tunnel to the frontend and serve its requests over it, so this worker needs no inbound port (see [Worker tunnels](#worker-tunnels)) | | `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS URL for backend installation and file staging | | `--nats-jwt` | `LOCALAI_NATS_JWT` | *(empty)* | Optional override for the `nats_jwt` returned at registration | | `--nats-user-seed` | `LOCALAI_NATS_USER_SEED` | *(empty)* | Optional override for `nats_user_seed` from registration | From 5108be222d8327fca9736ba9e331dbe48a06e4a3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 13:03:54 +0000 Subject: [PATCH 27/42] fix(worker): spec the tunnel's routing table, which was the SSRF boundary Review follow-up. One blocking finding and seven others. The blocking one first, and it is this project's recurring shape: the untested path. loopbackService is the function whose comment calls the discarded host "the security property this function exists for", and nothing tested it. The reviewer replaced its body with a dial of whatever the frontend named, no port range, and all 131 specs passed. Every spec installed the permissive test dialler, so the real routing table was exercised nowhere. It now has specs, and the property is stated as reachability rather than as a property of the code: a listener on 127.0.0.2 that only the frontend's target names must NOT be reached. Plus the port-range table, fixedService, loopbackAddr, tunnelEndpoint, and the table itself, which moved out of Run into tunnelServices so it can be built without starting a worker. One spec drives a real stream through that table over the wire, so the routing rules are exercised end to end at least once rather than only in isolation. The reviewer's mutation now reddens ten specs, and six narrower ones redden between two and four each, so no spec is riding on another. The shape changed too, not only the coverage. The dial address is built from a loopbackHost constant and strconv.Itoa of a validated int, so nothing derived from the wire reaches DialContext at all: restoring the hole takes ADDING a data flow, not deleting a check. And a taxonomy fix found while specifying it. A port outside this worker's allocator range was reported as unavailable, which tells a frontend to retry something that can never work. It is a bad request now, and a backend that is merely not listening yet stays unavailable, which is the retryable one. Agent nodes no longer get a tunnel credential. Nothing dials into an agent worker, so a tunnel replaces nothing for it and no client would open one, and the gate is at the mint site rather than in the handler: with no credential minted the hash stays empty and the existing empty-hash refusal covers it, so enforcement is structural. Two comments and one doc paragraph said an anonymous registrant gets a "working" credential. With auto-approve off the node is pending and the credential is inert, which is the distinction this same change argues three files away to justify minting for pending nodes at all. A refusal reason over the frame limit was cut on a byte boundary and could split a rune. It cuts on a rune boundary now, and the code survives truncation, which is what keeps a refusal classifiable. Also: the pending-node spec asserted only that a credential was non-empty, so a credential derived from the shared token passed it; it now pins per-node-ness the way the headline spec does. The tunnel handler's citations into nodes.go were stale before this branch landed, having been written against a file the same commit was editing, and are by function name now. The static-NATS path says plainly that an externally forced rotation locks it out until restart, and where that gets fixed. tunnelproto gained direct specs, including that a read failure is never reported as a refusal. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] --- core/http/app.go | 13 +- core/http/endpoints/cluster/connect.go | 9 +- core/http/endpoints/localai/nodes.go | 17 +- core/http/endpoints/localai/nodes_test.go | 40 +++- core/services/cluster/tunnelproto.go | 30 ++- core/services/cluster/tunnelproto_test.go | 186 +++++++++++++++ core/services/worker/tunnel.go | 73 +++++- core/services/worker/tunnel_test.go | 279 ++++++++++++++++++++++ core/services/worker/worker.go | 33 +-- docs/content/features/distributed-mode.md | 8 +- 10 files changed, 651 insertions(+), 37 deletions(-) create mode 100644 core/services/cluster/tunnelproto_test.go diff --git a/core/http/app.go b/core/http/app.go index fe790fb16e4e..5d3299027580 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -609,8 +609,17 @@ func API(application *application.Application) (*echo.Echo, error) { // credential at registration whether or not one is configured. What // is missing is the gate in FRONT of that. With no registration // token, RegisterNodeEndpoint validates nothing, so anyone who can - // reach this frontend can register a node and be handed a working - // tunnel credential for it. + // reach this frontend can register a node and be issued a tunnel + // credential for it. + // + // How far that gets them depends on the OTHER knob. With + // auto-approve on, the node is healthy at once and the credential + // works immediately. With it off, the node is pending, and the + // tunnel route refuses a pending node on every dial, so the + // credential is inert until an admin approves it and approval is + // the real gate. Worth stating precisely, because the same commit + // argues exactly this distinction three files away to justify + // minting for pending nodes at all. // // This warning replaced one that said the opposite, that tunnels // would refuse every dial without this token. That was true while diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go index 5e6ace97f7b6..8d175704669c 100644 --- a/core/http/endpoints/cluster/connect.go +++ b/core/http/endpoints/cluster/connect.go @@ -130,8 +130,13 @@ func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegi // Only StatusPending is refused. The rest of /api/node/ self-service // gates on nothing at all, but the two places that hand a node something // DURABLE both refuse a pending one: the agent worker's API key - // (core/http/endpoints/localai/nodes.go:224) and its NATS credential - // (nodes.go:293). A tunnel is that kind of grant, not a heartbeat: it is + // (provisionAgentWorkerKey, guarded at its call site in + // core/http/endpoints/localai/nodes.go) and its NATS credential + // (attachNatsJWT in the same file). Cited by NAME, not by line: the + // previous version of this comment cited line numbers into a file this + // same commit was editing, and both were stale before it landed. + // + // A tunnel is that kind of grant, not a heartbeat: it is // a standing pipe into the worker recorded in node_connections and // relayed to by every other replica. Draining and unhealthy nodes keep // their tunnels on purpose; draining means finish what you have, and a diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index 0be399a06f49..f674cb71a985 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -306,6 +306,21 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr // path in core/services/worker/worker.go does), because approval alone does not // prompt a re-registration and nothing else can hand it the secret. // +// Only BACKEND nodes get one, and that is a decision rather than an oversight. +// An agent worker serves no gRPC backends and no file staging; nothing dials +// into it at all, so a tunnel replaces nothing for it and there is no client on +// the agent side that would ever open one. Minting anyway would hand out a +// working credential for a pipe nobody drives, which is surface without a +// feature, and it would contradict every comment in this change that says +// "backend workers, the ones that tunnel". +// +// The gate lives HERE and not in ConnectHandler, which never looks at NodeType. +// It does not need to: an agent node's tunnel credential is never minted, so +// its TunnelTokenHash stays empty and the handler's empty-hash branch refuses +// it like any other node without one. Enforcement is therefore structural. The +// day agent workers want a tunnel, relaxing this condition is the whole change, +// and it has to be a deliberate one. +// // A failure to mint or to store is logged and the response goes out without the // token. Registration is what gets a worker into the cluster at all, and // failing it over a credential the worker does not need until it tunnels would @@ -313,7 +328,7 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr // tunnel_token, reports that it has no credential, and retries at its next // registration. func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) { - if node == nil { + if node == nil || node.NodeType != nodes.NodeTypeBackend { return } // crypto/rand.Text: at least 128 bits of randomness, no error to handle and diff --git a/core/http/endpoints/localai/nodes_test.go b/core/http/endpoints/localai/nodes_test.go index 5b2d6841462e..85a2600b6e88 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -154,15 +154,49 @@ var _ = Describe("Node HTTP handlers", func() { }) It("issues a tunnel credential to a node still awaiting approval", func() { - resp := register(`{"name":"worker-pending","address":"10.0.0.5:50051"}`, "", false) - Expect(resp["status"]).To(Equal(nodes.StatusPending)) // Deliberately unlike the agent API key and the NATS JWT, which are // both withheld from a pending node. Those work the moment they are // issued; this one does not, because the tunnel endpoint re-reads // the node's status on every dial and refuses a pending node. A // worker that registers exactly once would otherwise never receive // one, since approval alone prompts no re-registration. - Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + first := register(`{"name":"worker-pending","address":"10.0.0.5:50051","token":"shared"}`, "shared", false) + Expect(first["status"]).To(Equal(nodes.StatusPending)) + plaintext, _ := first["tunnel_token"].(string) + Expect(plaintext).ToNot(BeEmpty()) + + // Non-empty alone does not pin per-node-ness, and a review's + // variant of the "derived from the shared token" mutation stayed + // green on exactly that gap. A pending node's credential has to be + // as unpredictable and as per-node as an approved one's, since it + // becomes live the moment an admin approves. + Expect(plaintext).ToNot(Equal("shared")) + node, err := registry.Get(context.Background(), first["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext))) + Expect(node.TunnelTokenHash).ToNot(Equal(node.TokenHash)) + + second := register(`{"name":"worker-pending-2","address":"10.0.0.5:50052","token":"shared"}`, "shared", false) + Expect(second["status"]).To(Equal(nodes.StatusPending)) + Expect(second["tunnel_token"]).ToNot(Equal(plaintext)) + }) + + It("does not issue a tunnel credential to an agent node", func() { + // An agent worker serves no gRPC backends and no file staging; + // nothing dials into it, so a tunnel replaces nothing for it and no + // client on its side would open one. Minting anyway would be + // credential surface with no feature behind it. + // + // Enforcement is structural rather than a second check: with no + // credential minted, the node's hash stays empty and the tunnel + // route refuses it like any other node without one. + resp := register(`{"name":"agent-1","node_type":"agent"}`, "", true) + Expect(resp["node_type"]).To(Equal(nodes.NodeTypeAgent)) + Expect(resp).ToNot(HaveKey("tunnel_token")) + + node, err := registry.Get(context.Background(), resp["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(BeEmpty()) }) It("returns nats_jwt when account seed is configured", func() { diff --git a/core/services/cluster/tunnelproto.go b/core/services/cluster/tunnelproto.go index 39f5014b50b4..c6433459f78c 100644 --- a/core/services/cluster/tunnelproto.go +++ b/core/services/cluster/tunnelproto.go @@ -8,6 +8,7 @@ import ( "fmt" "io" "strings" + "unicode/utf8" ) // The framing every stream on a worker tunnel opens with. @@ -154,10 +155,7 @@ func WriteStreamRefusal(w io.Writer, reason error) error { }, reason.Error()) } frame := replyPrefixRefused + code + streamRequestSeparator + text - if len(frame) > maxTunnelFrame { - frame = frame[:maxTunnelFrame] - } - return writeFrame(w, frame) + return writeFrame(w, truncateRunes(frame, maxTunnelFrame)) } // ReadStreamReply reads the worker's answer. nil means the stream is now @@ -196,6 +194,30 @@ func ReadStreamReply(r io.Reader) error { } } +// truncateRunes cuts s to at most limit BYTES, on a rune boundary. +// +// A plain slice would cut mid-rune and put a lone continuation byte on the +// wire. Nothing breaks: the frame is length-prefixed so the framing survives, +// and the reader's string() tolerates invalid UTF-8. What it costs is the +// far side's log line ending in a replacement character, and a refusal reason +// exists to be read by a person, so it should not arrive damaged. +// +// The code that reaches this is always short; only a cause from a local service +// can be long enough to matter. +func truncateRunes(s string, limit int) string { + if len(s) <= limit { + return s + } + cut := limit + // utf8.RuneStart finds the first byte of a rune. Walking back from the + // limit lands on the start of the rune that would have been split, and at + // most 3 steps are needed since a UTF-8 rune is at most 4 bytes. + for cut > 0 && !utf8.RuneStart(s[cut]) { + cut-- + } + return s[:cut] +} + // writeFrame writes one length-prefixed frame in a single Write. // // One Write, not two: the underlying stream is a yamux stream whose writes diff --git a/core/services/cluster/tunnelproto_test.go b/core/services/cluster/tunnelproto_test.go new file mode 100644 index 000000000000..752fb9e9b84b --- /dev/null +++ b/core/services/cluster/tunnelproto_test.go @@ -0,0 +1,186 @@ +package cluster_test + +import ( + "bytes" + "encoding/binary" + "io" + "strings" + "unicode/utf8" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +var _ = Describe("Worker tunnel stream framing", func() { + Describe("the request frame", func() { + DescribeTable("round-trips a tag and a target", + func(tag, target string) { + var buf bytes.Buffer + Expect(cluster.WriteStreamRequest(&buf, tag, target)).To(Succeed()) + gotTag, gotTarget, err := cluster.ReadStreamRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(gotTag).To(Equal(tag)) + Expect(gotTarget).To(Equal(target)) + }, + Entry("a tag and an address", cluster.StreamTagGRPC, "127.0.0.1:50051"), + Entry("a tag with no target", cluster.StreamTagHTTP, ""), + // The split is on the FIRST separator, so a target containing one + // must survive intact. + Entry("a target containing a space", cluster.StreamTagGRPC, "a b c"), + ) + + It("consumes exactly the frame and not one byte of what follows", func() { + // Load-bearing: the stream is handed to gRPC or net/http right + // after this, and a reader that over-read would eat the start of + // their conversation. + var buf bytes.Buffer + Expect(cluster.WriteStreamRequest(&buf, cluster.StreamTagGRPC, "127.0.0.1:1")).To(Succeed()) + buf.WriteString("PRI * HTTP/2.0") + + _, _, err := cluster.ReadStreamRequest(&buf) + Expect(err).ToNot(HaveOccurred()) + rest, err := io.ReadAll(&buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(rest)).To(Equal("PRI * HTTP/2.0")) + }) + + DescribeTable("refuses a tag it could not encode unambiguously", + func(tag string) { + var buf bytes.Buffer + Expect(cluster.WriteStreamRequest(&buf, tag, "x")).ToNot(Succeed()) + Expect(buf.Len()).To(BeZero(), "a refused request must not put a partial frame on the wire") + }, + Entry("empty", ""), + // A tag with a space would silently move part of itself into the + // target, so it is refused at the writer rather than a round trip + // later. + Entry("containing a space", "grpc stream"), + ) + + It("refuses a frame that declares more than the limit without allocating it", func() { + var hdr [2]byte + binary.BigEndian.PutUint16(hdr[:], 65535) + _, _, err := cluster.ReadStreamRequest(bytes.NewReader(hdr[:])) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("over the")) + }) + + It("reports a truncated frame as a truncated read, not as a refusal", func() { + // ReadStreamRequest must never produce ErrStreamRequestInvalid: + // that sentinel is what a worker SENDS, and a reader producing it + // would leave a caller unable to tell "the peer refused me" from + // "I could not read the peer". + var hdr [2]byte + binary.BigEndian.PutUint16(hdr[:], 10) + _, _, err := cluster.ReadStreamRequest(bytes.NewReader(append(hdr[:], 'a'))) + Expect(err).To(MatchError(io.ErrUnexpectedEOF)) + Expect(err).ToNot(MatchError(cluster.ErrStreamRequestInvalid)) + }) + }) + + Describe("the reply frame", func() { + It("reads an acceptance as nil", func() { + var buf bytes.Buffer + Expect(cluster.WriteStreamAccepted(&buf)).To(Succeed()) + Expect(cluster.ReadStreamReply(&buf)).To(Succeed()) + }) + + DescribeTable("keeps the three refusals apart", + func(sent error, others []error) { + var buf bytes.Buffer + Expect(cluster.WriteStreamRefusal(&buf, sent)).To(Succeed()) + got := cluster.ReadStreamReply(&buf) + Expect(got).To(MatchError(sent)) + // The whole point. A caller gives up on an unknown tag, retries + // an unavailable target, and reports a bad request as its own + // bug; collapsing any pair makes one of those wrong. + for _, other := range others { + Expect(got).ToNot(MatchError(other)) + } + }, + Entry("unknown tag", cluster.ErrStreamTagUnknown, + []error{cluster.ErrStreamTargetUnavailable, cluster.ErrStreamRequestInvalid}), + Entry("unavailable target", cluster.ErrStreamTargetUnavailable, + []error{cluster.ErrStreamTagUnknown, cluster.ErrStreamRequestInvalid}), + Entry("invalid request", cluster.ErrStreamRequestInvalid, + []error{cluster.ErrStreamTagUnknown, cluster.ErrStreamTargetUnavailable}), + ) + + It("carries the reason text to the far side", func() { + var buf bytes.Buffer + Expect(cluster.WriteStreamRefusal(&buf, wrapReason(cluster.ErrStreamTagUnknown, "no-such-tag"))).To(Succeed()) + Expect(cluster.ReadStreamReply(&buf).Error()).To(ContainSubstring("no-such-tag")) + }) + + It("reports an unrecognised code as itself, not as the nearest known one", func() { + // A code from a newer worker. Mapping it onto a known sentinel + // would make a frontend retry forever against a refusal that means + // something else entirely. + var buf bytes.Buffer + writeRawFrame(&buf, "err teapot short and stout") + got := cluster.ReadStreamReply(&buf) + Expect(got).To(HaveOccurred()) + Expect(got.Error()).To(ContainSubstring("teapot")) + Expect(got).ToNot(MatchError(cluster.ErrStreamTagUnknown)) + Expect(got).ToNot(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(got).ToNot(MatchError(cluster.ErrStreamRequestInvalid)) + }) + + It("reports a failure to READ the reply as itself, never as a refusal", func() { + // A refusal proves the worker is connected and said no. A read + // failure means the tunnel broke. A caller that treated the second + // as the first would report a dead link as a policy decision. + got := cluster.ReadStreamReply(bytes.NewReader(nil)) + Expect(got).To(MatchError(io.EOF)) + Expect(got).ToNot(MatchError(cluster.ErrStreamTagUnknown)) + Expect(got).ToNot(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(got).ToNot(MatchError(cluster.ErrStreamRequestInvalid)) + }) + + It("truncates an over-long reason on a rune boundary, keeping it decodable", func() { + // Two-byte runes so a byte-boundary cut lands mid-rune for half of + // all lengths; the padding tunes the frame to land exactly there. + reason := wrapReason(cluster.ErrStreamTargetUnavailable, strings.Repeat("é", 2000)) + var buf bytes.Buffer + Expect(cluster.WriteStreamRefusal(&buf, reason)).To(Succeed()) + + got := cluster.ReadStreamReply(&buf) + Expect(got).To(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(utf8.ValidString(got.Error())).To(BeTrue(), + "the truncated reason reached the far side with a split rune in it") + }) + + It("still reports the code when the reason is truncated away", func() { + // The code must survive truncation: a refusal a frontend cannot + // classify is indistinguishable from a worker that hung up. + reason := wrapReason(cluster.ErrStreamTagUnknown, strings.Repeat("x", 4000)) + var buf bytes.Buffer + Expect(cluster.WriteStreamRefusal(&buf, reason)).To(Succeed()) + Expect(cluster.ReadStreamReply(&buf)).To(MatchError(cluster.ErrStreamTagUnknown)) + }) + }) +}) + +// wrapReason builds the shape the worker sends: a sentinel with a cause. +func wrapReason(sentinel error, text string) error { + return &reasonErr{sentinel: sentinel, text: text} +} + +type reasonErr struct { + sentinel error + text string +} + +func (e *reasonErr) Error() string { return e.sentinel.Error() + ": " + e.text } +func (e *reasonErr) Unwrap() error { return e.sentinel } + +// writeRawFrame puts a payload on the wire without going through the encoder, +// so a spec can present a frame the encoder would never produce. +func writeRawFrame(buf *bytes.Buffer, payload string) { + var hdr [2]byte + binary.BigEndian.PutUint16(hdr[:], uint16(len(payload))) + buf.Write(hdr[:]) + buf.WriteString(payload) +} diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index 5ab812d8c23e..469bd47c86a3 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -371,9 +371,7 @@ func (t *Tunnel) accept(ctx context.Context, stream net.Conn) (net.Conn, bool) { local, err := svc(ctx, target) if err != nil { - // An INFRASTRUCTURE failure, which a frontend may retry, and which must - // never be reported as the unknown tag above. - t.refuse(stream, fmt.Errorf("%w: %v", cluster.ErrStreamTargetUnavailable, err)) + t.refuse(stream, classifyServiceFailure(err)) return nil, false } @@ -389,6 +387,26 @@ func (t *Tunnel) accept(ctx context.Context, stream net.Conn) (net.Conn, bool) { return local, true } +// classifyServiceFailure decides which refusal a local service's error is. +// +// A service that has ALREADY classified its own failure keeps that +// classification. loopbackService does, and the distinction is not cosmetic: a +// target outside this worker's backend port range is a request this worker will +// never serve, while a backend that is not listening yet is a condition that +// clears on its own. Reporting the first as the second tells a frontend to +// retry something that can never work; reporting the second as the first makes +// it give up on a backend that is merely starting. +// +// Anything unclassified is infrastructure, because that is what an unadorned +// dial failure is, and it must never become the unknown-tag refusal: a tag this +// worker serves does not stop being served because one dial failed. +func classifyServiceFailure(err error) error { + if errors.Is(err, cluster.ErrStreamRequestInvalid) { + return err + } + return fmt.Errorf("%w: %v", cluster.ErrStreamTargetUnavailable, err) +} + // refuse reports why a stream will not be served and then ENDS it. // // The close is the part that matters and it is not optional. A worker that says @@ -567,22 +585,59 @@ func tunnelEndpoint(frontendURL, nodeID string) (string, error) { // // The port range is the one the worker's own port allocator hands to backend // processes, so a stream cannot be pointed at some unrelated service that -// happens to be listening on this host. +// happens to be listening on this host. It is only as tight as the allocator's +// range, which by default runs to 65535; a deployment that wants it narrow sets +// LOCALAI_GRPC_MAX_PORT, which narrows both at once. +// +// Note the SHAPE, not only the checks. Nothing derived from the wire reaches +// the dialler: the address is built from the loopbackHost constant and from +// strconv.Itoa of an int this function validated, so `target` itself has no +// path to DialContext at all. Relaxing this into an arbitrary-host dialler +// therefore takes ADDING a data flow rather than deleting a check, which is the +// difference between a guard and a property. It has specs either way; the shape +// is what stops a plausible refactor from quietly restoring the hole. 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("routing a tunnel stream: %q is not a host:port: %w", target, err) + return nil, fmt.Errorf("%w: routing a tunnel stream: %q is not a host:port: %v", + cluster.ErrStreamRequestInvalid, target, err) } port, err := strconv.Atoi(portStr) if err != nil { - return nil, fmt.Errorf("routing a tunnel stream: %q has no numeric port: %w", target, err) + return nil, fmt.Errorf("%w: routing a tunnel stream: %q has no numeric port: %v", + cluster.ErrStreamRequestInvalid, target, err) } if port < minPort || port > maxPort { - return nil, fmt.Errorf("routing a tunnel stream: port %d is outside this worker's backend range [%d, %d]", port, minPort, maxPort) + // Invalid rather than unavailable: no retry can bring a port + // outside this worker's own allocator range into it. + return nil, fmt.Errorf("%w: routing a tunnel stream: port %d is outside this worker's backend range [%d, %d]", + cluster.ErrStreamRequestInvalid, port, minPort, maxPort) } var d net.Dialer - return d.DialContext(ctx, "tcp", net.JoinHostPort("127.0.0.1", portStr)) + return d.DialContext(ctx, "tcp", net.JoinHostPort(loopbackHost, strconv.Itoa(port))) + } +} + +// loopbackHost is the only host any tunnel stream is ever dialled on. It is a +// constant so that "the worker dials itself and nothing else" is a fact about +// the code rather than a claim about its inputs. +const loopbackHost = "127.0.0.1" + +// tunnelServices builds the routing table the worker installs on its tunnel. +// +// It exists as its own function so the table can be specced. The table is the +// security boundary of this whole feature, and building it inline in Run left +// it reachable only by starting a worker, which meant it was covered by nothing +// and an arbitrary-host regression passed the entire suite. +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, and only within its own backend port + // range, is reachable through it. + cluster.StreamTagGRPC: loopbackService(basePort, cfg.effectiveMaxPort(basePort)), + cluster.StreamTagHTTP: fixedService(loopbackAddr(httpBindAddr)), } } @@ -611,7 +666,7 @@ func loopbackAddr(bindAddr string) string { return bindAddr } if host == "" || host == "0.0.0.0" || host == "::" || host == "[::]" { - return net.JoinHostPort("127.0.0.1", port) + return net.JoinHostPort(loopbackHost, port) } return bindAddr } diff --git a/core/services/worker/tunnel_test.go b/core/services/worker/tunnel_test.go index 09285c401591..b5724a542459 100644 --- a/core/services/worker/tunnel_test.go +++ b/core/services/worker/tunnel_test.go @@ -8,6 +8,7 @@ import ( "net" "net/http" "net/http/httptest" + "strconv" "strings" "sync/atomic" "time" @@ -213,6 +214,66 @@ var _ = Describe("Worker tunnel client", func() { Eventually(read, "10s").Should(Receive(BeNil())) Expect(string(buf)).To(Equal("ping")) }) + + It("routes through the worker's OWN table, ignoring the host the frontend names", func() { + // Every other spec in this file installs dialLocalTCP, which dials + // whatever it is handed. This one installs tunnelServices, the + // table Run installs, so the wire path is exercised against the + // real routing rules at least once. + backend := echoListenerOn("127.0.0.1:0") + 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 := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + // A host that is not this machine, and a port that is. + Expect(cluster.WriteStreamRequest(stream, cluster.StreamTagGRPC, + fmt.Sprintf("attacker.invalid:%d", port))).To(Succeed()) + + reply := awaitErr(func() error { return cluster.ReadStreamReply(stream) }) + Eventually(reply, "10s").Should(Receive(BeNil())) + + _, err = stream.Write([]byte("loopback")) + Expect(err).ToNot(HaveOccurred()) + buf := make([]byte, len("loopback")) + read := awaitErr(func() error { + _, err := io.ReadFull(stream, buf) + return err + }) + Eventually(read, "10s").Should(Receive(BeNil())) + Expect(string(buf)).To(Equal("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") + }) + + stream, err := session().OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(cluster.WriteStreamRequest(stream, cluster.StreamTagGRPC, "127.0.0.1:22")).To(Succeed()) + + reply := awaitErr(func() error { return cluster.ReadStreamReply(stream) }) + var got error + Eventually(reply, "10s").Should(Receive(&got)) + // Three refusals, three meanings. A frontend retries unavailable + // and gives up on this one. + Expect(got).To(MatchError(cluster.ErrStreamRequestInvalid)) + Expect(got).ToNot(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(got).ToNot(MatchError(cluster.ErrStreamTagUnknown)) + }) }) Describe("refusing a stream it cannot serve", func() { @@ -507,3 +568,221 @@ var _ = Describe("Worker tunnel client", func() { }) }) }) + +// echoListenerOn is echoListener bound to a specific address, so a spec can put +// a listener somewhere the worker must NOT reach. +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 +} + +// portOf returns the port a listener 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 +} + +// The routing table is the security boundary of the whole tunnel, and until now +// nothing exercised it: every spec above installs dialLocalTCP, which is exactly +// the permissive dialler loopbackService exists to prevent. A review turned +// loopbackService into an arbitrary-host dialler and all 131 specs passed. +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 := echoListenerOn("127.0.0.1:0") + DeferCleanup(func() { _ = ln.Close() }) + port := portOf(ln) + + conn, err := loopbackService(port, port)(ctx, ln.Addr().String()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + + _, err = conn.Write([]byte("hi")) + Expect(err).ToNot(HaveOccurred()) + buf := make([]byte, 2) + _, err = io.ReadFull(conn, buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("hi")) + }) + + It("ignores the host the frontend names and dials loopback anyway", func() { + ln := echoListenerOn("127.0.0.1:0") + DeferCleanup(func() { _ = ln.Close() }) + port := portOf(ln) + + // A host that is emphatically not this machine. If it were honoured + // the dial would fail or, far worse, succeed against something else. + conn, err := loopbackService(port, port)(ctx, fmt.Sprintf("attacker.invalid:%d", port)) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Expect(conn.RemoteAddr().String()).To(Equal(ln.Addr().String())) + }) + + It("does not reach a listener on another local address the frontend names", func() { + // The SSRF proof, stated as a reachability fact rather than as a + // property of the code. The only listener is on 127.0.0.2; nothing + // is on 127.0.0.1 at that port. A service that honoured the named + // host would connect; one that dials loopback cannot. + victim, err := net.Listen("tcp", "127.0.0.2:0") + if err != nil { + Skip("this host cannot bind a second loopback address: " + err.Error()) + } + DeferCleanup(func() { _ = victim.Close() }) + port := portOf(victim) + + conn, err := loopbackService(port, port)(ctx, victim.Addr().String()) + if err == nil { + _ = conn.Close() + Fail("the worker reached a host the frontend named, so a stream can steer it off loopback") + } + 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(HaveOccurred()) + // Invalid, not unavailable. No retry brings a port outside this + // worker's own allocator range into it, and telling a frontend + // to retry forever is how a refusal becomes a hang. + Expect(err).To(MatchError(cluster.ErrStreamRequestInvalid)) + Expect(err).ToNot(MatchError(cluster.ErrStreamTargetUnavailable)) + }, + 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 non-numeric port", "127.0.0.1:http", 50051, 50060), + Entry("no port at all", "127.0.0.1", 50051, 50060), + Entry("an empty target", "", 50051, 50060), + ) + + It("reports a backend that is not listening as unavailable, which a frontend may retry", func() { + // The other half of the taxonomy: a port IN range with nothing on + // it is a backend that has not started yet, not a bad request. + ln := echoListenerOn("127.0.0.1:0") + port := portOf(ln) + Expect(ln.Close()).To(Succeed()) + + _, err := loopbackService(port, port)(ctx, ln.Addr().String()) + Expect(err).To(HaveOccurred()) + Expect(classifyServiceFailure(err)).To(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(classifyServiceFailure(err)).ToNot(MatchError(cluster.ErrStreamRequestInvalid)) + }) + }) + + Describe("fixedService", func() { + It("reaches its own address whatever the frontend names", func() { + ln := echoListenerOn("127.0.0.1:0") + DeferCleanup(func() { _ = ln.Close() }) + + conn, err := fixedService(ln.Addr().String())(ctx, "attacker.invalid:9") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Expect(conn.RemoteAddr().String()).To(Equal(ln.Addr().String())) + }) + }) + + DescribeTable("loopbackAddr rewrites a bind address into a dialable one", + func(bind, want string) { + Expect(loopbackAddr(bind)).To(Equal(want)) + }, + // Dialling 0.0.0.0 only accidentally reaches localhost, and not on + // every platform, so the wildcard is replaced rather than dialled. + 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() { + // The table Run installs. Built by its own function precisely so this + // can be asserted without starting a worker. + It("serves exactly the two tags the frontend may name", func() { + cfg := &Config{ServeAddr: "0.0.0.0:50051"} + Expect(tunnelServices(cfg, "0.0.0.0:50050")).To(HaveLen(2)) + Expect(tunnelServices(cfg, "0.0.0.0:50050")).To(HaveKey(cluster.StreamTagGRPC)) + Expect(tunnelServices(cfg, "0.0.0.0:50050")).To(HaveKey(cluster.StreamTagHTTP)) + }) + + It("bounds the gRPC service by THIS worker's configured port range", func() { + cfg := &Config{ServeAddr: "0.0.0.0:50051", GRPCMaxPort: 50052} + svc := tunnelServices(cfg, "0.0.0.0:50050")[cluster.StreamTagGRPC] + + // The HTTP server's own port sits one below the base port, so a + // gRPC-tagged stream cannot be steered onto it. + _, err := svc(ctx, "127.0.0.1:50050") + Expect(err).To(MatchError(cluster.ErrStreamRequestInvalid)) + _, err = svc(ctx, "127.0.0.1:50053") + Expect(err).To(MatchError(cluster.ErrStreamRequestInvalid)) + }) + + // Pins that the HTTP service reaches the address Run configures. It + // does NOT pin the wildcard rewrite: on Linux dialling 0.0.0.0 reaches + // loopback anyway, so this spec stays green with loopbackAddr disabled. + // The loopbackAddr table above is what holds that, and it exists + // because the accident is not portable. + It("points the HTTP service at the worker's own server", func() { + ln := echoListenerOn("127.0.0.1:0") + DeferCleanup(func() { _ = ln.Close() }) + + cfg := &Config{ServeAddr: "0.0.0.0:50051"} + svc := tunnelServices(cfg, fmt.Sprintf("0.0.0.0:%d", portOf(ln)))[cluster.StreamTagHTTP] + + conn, err := svc(ctx, "ignored:1") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Expect(conn.RemoteAddr().String()).To(Equal(ln.Addr().String())) + }) + }) + + DescribeTable("tunnelEndpoint builds the URL the worker dials", + func(frontendURL, nodeID, want string) { + got, err := tunnelEndpoint(frontendURL, 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 passes through", "ws://frontend:8080", "n1", "ws://frontend:8080/api/cluster/connect?id=n1"), + Entry("wss passes through", "wss://frontend", "n1", "wss://frontend/api/cluster/connect?id=n1"), + // A frontend behind a path prefix keeps it: the path is appended, not + // assigned, exactly as the registration client builds its URLs. + Entry("a path prefix is kept", "https://host/localai", "n1", "wss://host/localai/api/cluster/connect?id=n1"), + Entry("a trailing slash 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 frontend URL it cannot dial", + func(frontendURL string) { + _, err := tunnelEndpoint(frontendURL, "n1") + Expect(err).To(HaveOccurred()) + }, + Entry("empty", ""), + // Refused rather than coerced: a worker silently dialling a scheme + // nobody configured is worse than one that says it cannot start. + Entry("a scheme that is not HTTP", "ftp://frontend"), + Entry("a bare host with no scheme", "frontend:8080/x"), + Entry("no host", "http://"), + ) +}) diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 829472ffb7e8..10b0b22c7209 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -16,7 +16,6 @@ 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" @@ -108,10 +107,20 @@ func Run(ctx *cliContext.Context, cfg *Config) error { } nodeID = res.ID // This path registers exactly once and never again, so the credential - // it holds cannot go stale by rotation from its own side. It can still - // be superseded from OUTSIDE, by a second worker registering under the - // same node name, and there is nothing this worker can do about that - // but log the 401 and keep retrying. + // it holds cannot go stale by rotation from its own side. + // + // It CAN be superseded from outside: Register upserts by NAME, so a + // second worker registering under this node's name rotates the row's + // credential, and this worker then fails every tunnel dial with 401 for + // the life of the process. It logs that once per backoff and never + // recovers on its own; a restart fixes it, because startup re-registers + // unconditionally. + // + // Deliberately NOT fixed here. Re-registering after repeated tunnel + // 401s is a decision about the worker's lifecycle, and it belongs with + // the change that removes this worker's inbound listeners, when a + // worker that cannot tunnel is a worker that cannot be reached at all. + // Today it can still be reached at the addresses it advertises. staticTunnelToken := res.TunnelToken tunnelToken = func() string { return staticTunnelToken } connectNats = func() (*messaging.Client, error) { @@ -191,19 +200,15 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // frontend URL or this node's identity is unusable, and a worker that // silently ran without its tunnel would look healthy while being // unreachable to everything that dials through it. - tunnelBasePort := cfg.effectiveBasePort() if cfg.WorkerTunnel { tunnel, terr := StartTunnel(shutdownCtx, TunnelConfig{ FrontendURL: cfg.RegisterTo, NodeID: nodeID, Token: tunnelToken, - Services: map[string]LocalService{ - // The frontend names a backend process by its port; the worker - // decides that only its own loopback, and only within its own - // backend port range, is reachable through it. - cluster.StreamTagGRPC: loopbackService(tunnelBasePort, cfg.effectiveMaxPort(tunnelBasePort)), - cluster.StreamTagHTTP: fixedService(loopbackAddr(httpAddr)), - }, + // Built by tunnelServices rather than inline, so the routing + // table, which is this feature's security boundary, is reachable + // from a spec without starting a worker. + Services: tunnelServices(cfg, httpAddr), }) if terr != nil { nodes.ShutdownFileTransferServer(httpServer) @@ -252,7 +257,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { }() // Process supervisor — manages multiple backend gRPC processes on different ports - basePort := tunnelBasePort + basePort := cfg.effectiveBasePort() // Buffered so NATS stop handler can send without blocking sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index e276acf406cf..22f61ba66025 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -116,7 +116,9 @@ A worker that presents a credential belonging to no node, or names a node ID the Unlike the agent worker's API key and its NATS credential, a tunnel credential **is** issued to a node still awaiting approval. It is inert until then: the tunnel route re-reads the node's status on every dial and refuses a pending one. Withholding it would instead strand workers that register exactly once, since approval on its own prompts no re-registration. -A tunnel credential does not replace `LOCALAI_REGISTRATION_TOKEN`. Without one, node registration itself is unauthenticated, so anyone who can reach the frontend can register a worker and be issued a working tunnel credential for it. LocalAI warns about that at startup. +A tunnel credential does not replace `LOCALAI_REGISTRATION_TOKEN`. Without one, node registration itself is unauthenticated, so anyone who can reach the frontend can register a worker and be issued a tunnel credential for it. How far that gets them depends on auto-approve: with auto-approve on the node is healthy at once and the credential works immediately; with it off the node is pending and the credential is inert until an admin approves, so approval is the real gate. LocalAI warns about the missing token at startup. + +Only **backend** nodes are issued one. An agent worker has no inbound surface for the tunnel to replace and no client for it, so minting one would widen the credential surface for nothing; its row keeps an empty tunnel credential and the tunnel route refuses it like any other node without one. The tunnel lands on exactly one frontend replica, and that replica records itself as the owner of the worker's connection in the `node_connections` table. When the socket dies the claim is dropped with it. If the replica stalls long enough for its peers to reap it, it re-claims the tunnels it still holds on a live session as soon as it re-registers, skipping any whose socket has already gone. That re-claim needs the replica to have an advertised address: without one it never had an instance row to begin with, and its tunnels are usable only by the replica holding them. @@ -135,7 +137,9 @@ The worker holds the tunnel with one goroutine: it dials, serves the frontend's | `grpc` | a backend process on this worker | the port; the host is discarded and only `127.0.0.1` is dialled, within the worker's own backend port range | | `http` | the worker's own file-transfer and backend-log server | ignored; there is one such server and only the worker knows where it bound | -A stream naming a tag the worker does not serve, or a local service it could not reach, is refused with a reason and the stream is **ended** rather than left open. Those two refusals are distinct on the wire on purpose: a frontend gives up on the first and may retry the second. One bad stream never affects the others or the session. +The `grpc` row is the security boundary of the tunnel, and it is worth being explicit about it. A tunnel terminates inside the worker process, so a stream arriving on it can reach anything the worker can reach; if the frontend could name the host, whoever holds the frontend end could make every worker in the fleet dial arbitrary addresses on its private network. The worker therefore builds the dial address from a constant `127.0.0.1` and a port it has validated, and the string from the wire never reaches the dialler at all. The port range is the one the worker's own allocator hands to backend processes, which by default runs to 65535; setting `LOCALAI_GRPC_MAX_PORT` narrows the allocator and this range together, and a worker with a known backend count should set it. + +A stream naming a tag the worker does not serve, a target outside that port range, or a local service it could not reach, is refused with a reason and the stream is **ended** rather than left open. Those refusals are distinct on the wire on purpose: an unknown tag and an out-of-range target are requests this worker will never serve, while an unreachable local service is a backend that has not started yet. A frontend gives up on the first two and may retry the third. One bad stream never affects the others or the session. Reconnects use exponential backoff with jitter: the interval doubles from 500ms up to a ceiling of 30 seconds, and each wait is drawn between half of that interval and all of it, so no worker ever spins and a fleet that lost the same replica does not come back in lockstep. The interval returns to its floor only after a session that lasted at least 30 seconds. That last part is what stops a rolling frontend restart, where every dial succeeds and then dies moments later, from turning a fleet of workers into a retry storm against the first replica back up. A worker that is refused (`401`, `403`) keeps retrying on the same schedule rather than exiting: a re-registration or an admin approval fixes both without restarting it. From 3b6d32c1c4b99a294f98cf93e2d84d4c52ea25e3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 13:33:07 +0000 Subject: [PATCH 28/42] fix(worker): make the tunnel credential's node-type gate actually structural Re-review follow-up, three items. Two are the overclaiming-comment class again, and the first is that class with a real defect underneath it. attachTunnelToken said "enforcement is therefore structural": an ineligible node never gets a credential, so its hash stays empty and the tunnel route's empty-hash branch does the refusing. That was true for a node that had always been an agent and false for one that had not. Register upserts by NAME, so a backend node re-registering as an agent keeps its ID, and Register's struct Updates zero-skips the credential column while writing the new node_type. The early return left the credential the node earned as a backend sitting on a row that is now an agent, and ConnectHandler never looks at node_type. Fixed by making the claim true rather than by softening it, because the mint-site gate was chosen precisely on the grounds that it was structural: an ineligible node now has its column CLEARED, unconditionally, so the invariant does not depend on what the row happened to contain. A spec pins it and was red before the change. Same shape as the Register-upserts-by-name hazard already carried forward: a name is not an identity. Second, loopbackHost claimed to be the only host any tunnel stream is ever dialled on. It is not: fixedService dials whatever Run built it from, which is this worker's own LOCALAI_HTTP_ADDR, and loopbackAddr rewrites only a wildcard bind, so an operator who binds the file-transfer server to a routable address gets a routable dial. The property that matters is narrower and is what the comment says now: the frontend cannot STEER the dial. The grpc tag builds its address from a constant and a validated port with nothing from the wire reaching the dialler, and the http tag ignores its target entirely. Worth stating exactly rather than summarising, because the argument about what a stream can reach rests on knowing which hosts are reachable, and an overstatement at that site is what would let someone conclude the constant alone is doing the work. Third, a spec named "without allocating it" measured no allocation. It now asserts the mechanism the defence actually rests on, that the reader consumes the two length bytes and not one byte of the body, through a counting reader. The input carries a body on purpose: against input that ends after the header the assertion would pass with the limit check deleted. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] --- core/http/endpoints/localai/nodes.go | 32 ++++++++++++++++++----- core/http/endpoints/localai/nodes_test.go | 28 ++++++++++++++++++++ core/services/cluster/tunnelproto_test.go | 30 +++++++++++++++++++-- core/services/worker/tunnel.go | 25 +++++++++++++++--- 4 files changed, 104 insertions(+), 11 deletions(-) diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index f674cb71a985..0d3f3ebb3f92 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -315,11 +315,21 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr // "backend workers, the ones that tunnel". // // The gate lives HERE and not in ConnectHandler, which never looks at NodeType. -// It does not need to: an agent node's tunnel credential is never minted, so -// its TunnelTokenHash stays empty and the handler's empty-hash branch refuses -// it like any other node without one. Enforcement is therefore structural. The -// day agent workers want a tunnel, relaxing this condition is the whole change, -// and it has to be a deliberate one. +// It does not need to, PROVIDED an ineligible node ends up with no credential +// rather than merely being handed no new one, because the handler's empty-hash +// branch is what does the refusing. So this CLEARS the column instead of +// returning early, and the difference is not theoretical: Register upserts by +// NAME, so a backend node re-registering as an agent keeps its ID, and +// Register's struct Updates zero-skips TunnelTokenHash while writing the new +// node_type. An early return left a live credential on a row that had become an +// agent. Clearing is what makes "enforcement is structural" true. +// +// It clears unconditionally rather than only when something is there, so the +// invariant holds without depending on what the row happened to contain. The +// cost is one UPDATE per agent registration. +// +// The day agent workers want a tunnel, relaxing the eligibility condition is +// the whole change, and it has to be a deliberate one. // // A failure to mint or to store is logged and the response goes out without the // token. Registration is what gets a worker into the cluster at all, and @@ -328,7 +338,17 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr // tunnel_token, reports that it has no credential, and retries at its next // registration. func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) { - if node == nil || node.NodeType != nodes.NodeTypeBackend { + if node == nil { + return + } + if node.NodeType != nodes.NodeTypeBackend { + // Cleared, not skipped. SetTunnelTokenHash writes the single column + // directly rather than through a struct update, so unlike Register it + // can write an empty value; see its doc. + if err := registry.SetTunnelTokenHash(ctx, node.ID, ""); err != nil { + xlog.Error("Failed to clear the tunnel credential of a node that is not a backend worker", + "node", node.Name, "type", node.NodeType, "error", err) + } return } // crypto/rand.Text: at least 128 bits of randomness, no error to handle and diff --git a/core/http/endpoints/localai/nodes_test.go b/core/http/endpoints/localai/nodes_test.go index 85a2600b6e88..dababff38420 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -199,6 +199,34 @@ var _ = Describe("Node HTTP handlers", func() { Expect(node.TunnelTokenHash).To(BeEmpty()) }) + It("clears a tunnel credential when a node stops being a backend node", func() { + // Register upserts BY NAME, so a node can change node_type in place. + // Skipping the mint on the way through leaves the credential the + // node earned as a backend sitting on a row that is now an agent: + // Register's struct Updates zero-skips the column while writing the + // new node_type, so nothing else clears it. ConnectHandler never + // looks at node_type, so that stale hash is a usable tunnel + // credential for a node type that is not supposed to hold one. + // + // This is the same shape as the Register-upserts-by-name hazard + // already carried forward: a name is not an identity. + backend := register(`{"name":"shifty","address":"10.0.0.7:50051"}`, "", true) + Expect(backend["tunnel_token"]).ToNot(BeEmpty()) + + agent := register(`{"name":"shifty","node_type":"agent"}`, "", true) + Expect(agent["id"]).To(Equal(backend["id"]), "re-registration must keep the node identity") + Expect(agent["node_type"]).To(Equal(nodes.NodeTypeAgent)) + Expect(agent).ToNot(HaveKey("tunnel_token")) + + node, err := registry.Get(context.Background(), backend["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + // The claim the gate makes is that an ineligible node HAS no + // credential, not merely that it was not handed a new one. Only + // then is the empty-hash refusal in ConnectHandler the enforcement. + Expect(node.TunnelTokenHash).To(BeEmpty(), + "the node kept the credential it earned as a backend, so the mint-site gate is not structural") + }) + It("returns nats_jwt when account seed is configured", func() { akp, err := nkeys.CreateAccount() Expect(err).ToNot(HaveOccurred()) diff --git a/core/services/cluster/tunnelproto_test.go b/core/services/cluster/tunnelproto_test.go index 752fb9e9b84b..2664d9d509f5 100644 --- a/core/services/cluster/tunnelproto_test.go +++ b/core/services/cluster/tunnelproto_test.go @@ -59,12 +59,25 @@ var _ = Describe("Worker tunnel stream framing", func() { Entry("containing a space", "grpc stream"), ) - It("refuses a frame that declares more than the limit without allocating it", func() { + It("refuses an over-long declared length after reading only the header", func() { + // The name used to say "without allocating it" and the spec + // measured nothing of the sort. What is actually checkable, and is + // the mechanism the defence rests on, is that the reader STOPS: it + // consumes the two length bytes and not one byte of the body, so a + // peer cannot make it allocate or read on demand. + // + // The body is present in the input on purpose. With an input that + // ends after the header, a reader that went on to read the body + // would still consume nothing more, and this assertion would pass + // with the limit check deleted. var hdr [2]byte binary.BigEndian.PutUint16(hdr[:], 65535) - _, _, err := cluster.ReadStreamRequest(bytes.NewReader(hdr[:])) + src := &countingReader{r: bytes.NewReader(append(hdr[:], make([]byte, 4096)...))} + + _, _, err := cluster.ReadStreamRequest(src) Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("over the")) + Expect(src.n).To(Equal(2), "the reader consumed part of a frame it had already refused") }) It("reports a truncated frame as a truncated read, not as a refusal", func() { @@ -176,6 +189,19 @@ type reasonErr struct { func (e *reasonErr) Error() string { return e.sentinel.Error() + ": " + e.text } func (e *reasonErr) Unwrap() error { return e.sentinel } +// countingReader records how many bytes were actually consumed, so a spec can +// assert where a reader stopped rather than only what it returned. +type countingReader struct { + r io.Reader + n int +} + +func (c *countingReader) Read(p []byte) (int, error) { + n, err := c.r.Read(p) + c.n += n + return n, err +} + // writeRawFrame puts a payload on the wire without going through the encoder, // so a spec can present a frame the encoder would never produce. func writeRawFrame(buf *bytes.Buffer, payload string) { diff --git a/core/services/worker/tunnel.go b/core/services/worker/tunnel.go index 469bd47c86a3..0eabacb0f9e3 100644 --- a/core/services/worker/tunnel.go +++ b/core/services/worker/tunnel.go @@ -619,9 +619,28 @@ func loopbackService(minPort, maxPort int) LocalService { } } -// loopbackHost is the only host any tunnel stream is ever dialled on. It is a -// constant so that "the worker dials itself and nothing else" is a fact about -// the code rather than a claim about its inputs. +// loopbackHost is the host every stream the FRONTEND CAN STEER is dialled on. +// +// It is a constant so that "a stream cannot choose where the worker dials" is a +// fact about the code rather than a claim about its inputs: the grpc tag builds +// its address from this and a port it validated, and nothing derived from the +// wire reaches the dialler. +// +// It is NOT the only host this file ever dials, and the difference is worth +// stating exactly rather than summarising, because the whole argument about +// what a stream can reach rests on knowing which hosts are reachable, and an +// overstatement here is what would let a future reader conclude the constant +// alone is doing the work. +// +// fixedService dials whatever address it was constructed with. Run constructs +// it from this worker's own LOCALAI_HTTP_ADDR, which an operator may set to a +// routable address; loopbackAddr only rewrites a WILDCARD bind, and leaves an +// explicit host alone on purpose, because a server bound to one address is not +// reachable on another. So the http tag can dial a non-loopback host. That host +// is one the OPERATOR configured for this worker's own server, never one a +// stream names: fixedService ignores its target entirely. The property 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 the worker installs on its tunnel. From cce914b3fdd343042117168e5883ef1d575fb590 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 14:12:17 +0000 Subject: [PATCH 29/42] feat(cluster): relay a peer's stream onto a worker tunnel held here A worker holds ONE tunnel and it lands on ONE frontend replica, so with N replicas behind a load balancer roughly (N-1)/N of requests arrive somewhere that cannot reach the worker directly. This is the piece that carries them: the SessionStore stream handler reads which worker a peer's stream is for, opens a stream on the tunnel this replica holds, and splices the two. Splice has had no production caller since phase 1. It has one now, and being the first caller it settles the two endings phase 1 deliberately left open, both of which read as normal termination until now: - a peer-initiated *StreamError{Remote: true}, which yamux builds only from an RST frame the far side sent (stream.go:432-449); a reset this side asks for carries Remote: false, and Splice never resets anything, its own Close sending a FIN; - a graceful ErrRemoteGoAway, which handleGoAway returns for code goAwayNormal (session.go:829-833) and close hands unwrapped to every live stream (session.go:328-337). Both truncate whatever was in flight. Reporting them as normal termination is how a half-finished inference comes to look like a short one that completed, so both are now reported; the local forms stay silent, because those are the teardown Splice provokes itself. The decision cannot live in a caller reading Splice's result, since a result already mapped to nil carries nothing left to reclassify, so it lives at the classifier with the reasoning beside it. The relay logs it at debug: a client cancelling a relayed request produces one per cancellation, and the truncation is separately visible to the frontend's own gRPC or HTTP client. The relay hop gets its own request and reply frames. They have to be distinct from the worker tunnel's, because a relayed stream carries both hops' frames back to back, and a vocabulary shared between them would let a reader applied to the wrong hop hand back a plausible sentinel belonging to the other. Its three refusals stay apart for the reason the worker's three do: ErrNotOwner is a routing fact and the caller should resolve the owner again; unavailable is infrastructure at this replica and a retry is worth something; bad-request is the caller's bug. None of them is, or may be built over, an absence error. One hop, always. A stream naming a worker this replica does not hold is refused, never resolved and relayed onward, so a stale ownership row cannot become a loop between two replicas each certain the other holds the worker. PeerPool is constructed and closed alongside SessionStore, so both halves of the peer mesh now have an owner and a shutdown. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 33 +- core/services/cluster/peerlink.go | 6 +- core/services/cluster/relay.go | 333 +++++++++++++++++++ core/services/cluster/relay_internal_test.go | 63 ++++ core/services/cluster/relay_test.go | 291 ++++++++++++++++ core/services/cluster/sessions.go | 4 + core/services/cluster/splice.go | 118 +++++-- core/services/cluster/splice_test.go | 16 +- 8 files changed, 819 insertions(+), 45 deletions(-) create mode 100644 core/services/cluster/relay.go create mode 100644 core/services/cluster/relay_internal_test.go create mode 100644 core/services/cluster/relay_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 30b297913f06..074d261d912b 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -54,8 +54,14 @@ type DistributedServices struct { // peer-reachable address could be determined, which leaves this replica // invisible to its peers but otherwise fully functional. Membership *cluster.Membership - // PeerSessions owns the peer links other replicas dialled into this one. + // PeerSessions owns the peer links other replicas dialled into this one, + // and relays the streams that arrive on them onto the worker tunnels this + // replica holds. PeerSessions *cluster.SessionStore + // Peers owns the peer links this replica dialled OUT, the mirror of + // PeerSessions. It is what the relaying dialer opens a stream on when a + // request arrives here for a worker another replica holds. + Peers *cluster.PeerPool // Tunnels holds the worker tunnels this replica has accepted and keeps the // node_connections table agreeing with them. It is handed to the membership // loop, which re-claims what it holds after this replica has been reaped, @@ -81,6 +87,13 @@ func (ds *DistributedServices) Shutdown() { if ds.PeerSessions != nil { ds.PeerSessions.CloseAll() } + // Both halves of the peer mesh go down together. A pool left open + // holds a WebSocket and two yamux loop goroutines per peer for as long + // as the process lives, and an Open after this reports ErrPoolClosed, + // which is a fact about this process and never node absence. + if ds.Peers != nil { + ds.Peers.Close() + } if ds.Health != nil { ds.Health.Stop() } @@ -193,10 +206,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Replica membership. NewNodeRegistry has just migrated the tables this // reads, so it has to come after it. clusterRegistry := cluster.NewRegistry(authDB) - // Accepted peer links are held with no stream handler: this replica has - // somewhere to put a link a peer dials, and refuses the streams on it, - // because nothing relays worker traffic yet. - peerSessions := cluster.NewSessionStore(nil) var membership *cluster.Membership if advertised, err := advertisedPeerAddr(cfg); err != nil { // Not fatal. A replica that cannot publish an address still serves @@ -233,6 +242,19 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade membership.SetTunnels(tunnels) } + // The links peers dial IN, with the relay installed on them. This is what + // makes more than one replica work: a worker holds one tunnel, it lands on + // one replica, and every request that arrives anywhere else reaches the + // worker through this handler. Passing nil here would leave every such + // request refused, promptly and only at debug level, which presents as a + // worker that is connected and unusable from most of the deployment. + peerSessions := cluster.NewSessionStore(cluster.NewRelay(tunnels).Stream) + // The links this replica dials OUT, the other half of the same mesh. It + // authenticates with the registration token because that is the token the + // peer route checks (see RegisterClusterRoutes); two different tokens here + // would make every peer dial 401 with nothing naming the mismatch. + peers := cluster.NewPeerPool(cfg.Distributed.InstanceID, cfg.Distributed.RegistrationToken, clusterRegistry) + // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an // operator can pin placement to a stable name like "production" and have it @@ -524,6 +546,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Cluster: clusterRegistry, Membership: membership, PeerSessions: peerSessions, + Peers: peers, Tunnels: tunnels, }, nil } diff --git a/core/services/cluster/peerlink.go b/core/services/cluster/peerlink.go index 961c1a3fa3d9..047fcbc5472a 100644 --- a/core/services/cluster/peerlink.go +++ b/core/services/cluster/peerlink.go @@ -233,8 +233,10 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { // has left the deployment but is still listening keeps a live WebSocket and the // two yamux loop goroutines behind it for as long as this process runs; a peer // that is genuinely gone is reclaimed by the 30s keepalive default, so the real -// exposure is narrow. There is no Forget because nothing yet knows which peers -// have left; Task 5's ownership work is where that knowledge appears. +// exposure is narrow. There is no Forget: the membership sweep does know which +// replicas have left, but nothing plumbs that knowledge to this pool, so an +// entry for a departed peer outlives it and only the keepalive reclaims what it +// holds. func (p *PeerPool) link(peerID string) (*peerLink, error) { p.mu.Lock() defer p.mu.Unlock() diff --git a/core/services/cluster/relay.go b/core/services/cluster/relay.go new file mode 100644 index 000000000000..99164c6fbd74 --- /dev/null +++ b/core/services/cluster/relay.go @@ -0,0 +1,333 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "cmp" + "context" + "errors" + "fmt" + "io" + "net" + "strings" + "time" + + "github.com/mudler/xlog" +) + +// The framing every stream on a PEER link opens with. +// +// A worker holds one tunnel and it lands on one frontend replica, so every +// other replica reaches that worker by relaying through the one that holds it. +// The peer link carries traffic for every worker its far side owns, so a stream +// on it means nothing until it says which worker it is for; that is this frame. +// +// A relayed stream therefore carries TWO request frames back to back: this one, +// which the owning replica consumes, and the worker tunnel's own (tunnelproto) +// frame, which crosses untouched and is answered by the worker. A dialler reads +// one reply from each, in that order. +// +// The two vocabularies are deliberately disjoint. "relay-ok" is not "ok", and +// none of the three refusal codes below is spelled like a tunnel code, so a +// reader applied to the wrong hop fails with "unrecognised reply" instead of +// handing back a plausible sentinel that belongs to the other hop. Getting that +// wrong would report a worker's refusal as the owning replica's, and a caller +// would retry against the wrong end of the path. +const ( + relayReplyAccepted = "relay-ok" + relayCodeNotOwner = "relay-not-owner" + relayCodeUnavailable = "relay-unavailable" + relayCodeBadRequest = "relay-bad-request" +) + +// The refusals the relay hop can send, beyond ErrNotOwner which it shares with +// the local path. +// +// Three conditions, kept apart, for the reason the worker's three are kept +// apart. ErrNotOwner is a ROUTING fact: the worker may be perfectly healthy on +// another replica, and the caller should resolve the owner again. This one is +// INFRASTRUCTURE at the owning replica: the tunnel is held right here and its +// session will not carry a stream, so a retry is worth something and looking +// elsewhere is not. ErrRelayRequestInvalid is the CALLER's bug and no retry +// helps. +// +// None of them is, or may ever be built over, an absence error. A refusal is +// proof that a replica answered, and reporting absence would tell a scheduler +// that a worker which is connected has gone away. +var ( + ErrRelayUnavailable = errors.New("cluster: the owning replica could not open a stream to that worker") + ErrRelayRequestInvalid = errors.New("cluster: the owning replica rejected the relay request as malformed") +) + +// WriteRelayRequest names the worker a peer stream is for. +// +// An empty node id is refused here rather than on the wire, so a caller with a +// bug learns at once instead of a round trip later. +func WriteRelayRequest(w io.Writer, nodeID string) error { + if nodeID == "" { + return fmt.Errorf("writing a relay request: empty node id") + } + return writeFrame(w, nodeID) +} + +// ReadRelayRequest reads the opening frame of a peer stream. +// +// A malformed frame is an ordinary error, NOT ErrRelayRequestInvalid: that +// sentinel is what a relay SENDS to describe a refusal, and producing it here +// would leave a caller unable to tell "the peer refused my request" from "I +// could not read the peer's". +func ReadRelayRequest(r io.Reader) (string, error) { + payload, err := readFrame(r) + if err != nil { + return "", fmt.Errorf("reading a relay request: %w", err) + } + if payload == "" { + // An empty payload is a well-formed frame naming no worker. Treating + // it as a node called "" would send the caller a routing refusal for a + // request no replica can ever serve, so it stays the caller's bug. + return "", fmt.Errorf("reading a relay request: empty node id") + } + return payload, nil +} + +// WriteRelayAccepted tells the peer the stream now carries the worker tunnel's +// own conversation. Everything after this frame belongs to that hop. +func WriteRelayAccepted(w io.Writer) error { return writeFrame(w, relayReplyAccepted) } + +// WriteRelayRefusal reports why a peer's stream will not be relayed. The caller +// closes the stream afterwards; this only says why. +// +// An unclassified reason is sent as bad-request with its text attached, rather +// than dropped: a refusal a peer cannot read is indistinguishable from a +// replica that hung up, and those are different problems. +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 { + // The frame is length-prefixed so a newline would not corrupt it, + // but this text lands in a log line on the far side, and a cause + // spanning lines is what makes one unsearchable. + if r == '\n' || r == '\r' { + return ' ' + } + return r + }, reason.Error()) + } + frame := replyPrefixRefused + code + streamRequestSeparator + text + return writeFrame(w, truncateRunes(frame, maxTunnelFrame)) +} + +// ReadRelayReply reads the owning replica's answer to a relay request. nil +// means the stream is now the worker tunnel's. +// +// A failure to READ the reply is returned as itself and never as one of the +// refusal sentinels: a refusal means a replica answered, a read failure means +// the peer link broke, and reporting the second as the first would present a +// dead link as a policy 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. It is a routing fact, and the sentinels + // it must never be confused with are ErrNoConnection (the worker is + // connected nowhere) and ErrPeerUnreachable (a replica will not + // answer): a caller acts on those by giving up on the worker or by + // retrying the peer, and on this one by resolving the owner again. + 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) + default: + // A code from a newer replica. Carried out as-is rather than mapped + // onto the nearest known one, so a caller does not retry forever + // against a refusal that means something else entirely. + return fmt.Errorf("relay stream refused with unrecognised code %q: %s", code, text) + } +} + +const ( + // relayHeaderTimeout bounds how long a peer stream may go without naming + // the worker it is for. Without it, a dialler killed between OpenStream and + // its first write holds a relay goroutine and a stream slot until the whole + // peer link dies, which is minutes on the default keepalive. + relayHeaderTimeout = 15 * time.Second + + // relayOpenTimeout bounds opening the worker-side stream. yamux blocks an + // Open once AcceptBacklog SYNs are in flight, waiting on synCh rather than + // failing (go-yamux/v5@v5.1.0/session.go:205-212); it honours the context, + // which is the only reason there is one here. Without the bound, a worker + // that has stopped accepting would turn a refusable condition into a parked + // peer, which is the one outcome this path exists to avoid. + relayOpenTimeout = 15 * time.Second +) + +// Relay splices a stream a peer opened onto a worker tunnel this replica holds. +// +// It is the piece that makes more than one frontend replica work at all: a +// worker holds ONE tunnel, it lands on ONE replica, and with N replicas behind +// a load balancer roughly (N-1)/N of requests arrive somewhere else. Those +// requests reach the worker through here. +// +// One hop, always. A stream naming a worker this replica does not hold is +// refused, never resolved and relayed onward. A second hop would turn a stale +// ownership row into a loop between two replicas, each certain the other holds +// the worker, and the loop would carry the caller's request around it; the +// dialling replica re-resolving the owner is both cheaper and terminating. +type Relay struct { + tunnels *TunnelRegistry + + // Timeouts are fields rather than constants read directly so a spec can + // exercise the deadline without waiting out a production value. They are + // not operator knobs and are not plumbed to configuration. + headerTimeout time.Duration + openTimeout time.Duration +} + +// NewRelay returns the relay for the tunnels this replica holds. Its Stream +// method is the SessionStore stream handler. +func NewRelay(tunnels *TunnelRegistry) *Relay { return newRelay(tunnels, 0, 0) } + +func newRelay(tunnels *TunnelRegistry, 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 that stream on every path. +func (r *Relay) Stream(peerID string, stream net.Conn) { + // SessionStore runs this on a bare goroutine, so an unrecovered panic here + // ends the PROCESS, taking down every other replica's traffic through this + // one. It covers what runs on this goroutine: the frame read, the registry + // lookup and the open. It cannot cover a panic inside Splice's own copy + // goroutines, and it deliberately does not re-panic, because there is no + // recovery middleware above a goroutine the HTTP layer has already + // returned from. + 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 { + // accept has already answered and closed the stream. + return + } + + // Splice owns closing both ends from here. + // + // The error is logged at DEBUG and nowhere else. Every relayed request that + // a client abandons mid-stream produces one, so a warning here would be one + // line per cancelled inference; and the failures that are not cancellations + // are already visible to the frontend, whose gRPC or HTTP client sees a + // response that ended without its trailers or its final chunk. What this + // line adds is the only view from the middle of the path: which node, on + // which peer link, and what yamux actually said. + 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 worker-side stream. +// The second result is false when the stream was refused, in which case the +// refusal has been 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 { + // Nothing is readable on a stream whose deadline cannot be set, so this + // is reported as infrastructure rather than pushed past. + r.refuse(peerID, stream, fmt.Errorf("%w: arming the request deadline: %v", ErrRelayUnavailable, err)) + return nil, false + } + + nodeID, err := ReadRelayRequest(stream) + if err != nil { + // Includes the deadline above expiring. Both are "this stream never + // said which worker it wanted", which is the dialling replica's bug + // and not something a retry against this one resolves. + r.refuse(peerID, stream, fmt.Errorf("%w: %v", ErrRelayRequestInvalid, err)) + return nil, false + } + + // Cleared before the open rather than after the reply: everything past this + // frame belongs to the worker tunnel's conversation, which brings its own + // deadlines, and one left armed here would abort a long inference stream in + // the middle. + 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 + } + + // Not the peer's deadline, because there is none to inherit: a yamux stream + // carries no context. This bounds only the open, so a request that gets + // past it is never cut short by it. + ctx, cancel := context.WithTimeout(context.Background(), r.openTimeout) + defer cancel() + + local, err := r.tunnels.Open(ctx, nodeID) + if err != nil { + if errors.Is(err, ErrNotOwner) { + // Passed through as itself. The worker is very likely connected and + // healthy somewhere else, and this is the one answer that tells the + // caller to look for it there. + r.refuse(peerID, stream, err) + return nil, false + } + // Everything else is this replica failing, and it must NOT become + // ErrNotOwner. The tunnel is held right here, so sending the caller + // looking elsewhere would send it back to this same replica; and it + // must not become absence either, because the worker is attached and a + // scheduler told otherwise would reclaim what it is running. + r.refuse(peerID, stream, fmt.Errorf("%w: %v", ErrRelayUnavailable, err)) + return nil, false + } + + if err := WriteRelayAccepted(stream); err != nil { + // The peer never learns the stream was accepted, so it cannot be used. + // Closing the worker-side stream here is what stops one leaking per + // 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 will not be relayed and then ENDS it. +// +// The close is the part that matters and it 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, which reads as a slow replica rather than a refused +// request, and no deadline on the far side can tell those apart. The reply is +// what makes the refusal legible; the close is what makes it prompt. +// +// The reply is therefore 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/cluster/relay_internal_test.go b/core/services/cluster/relay_internal_test.go new file mode 100644 index 000000000000..855ac58653bf --- /dev/null +++ b/core/services/cluster/relay_internal_test.go @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "net" + "time" + + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// This spec is in-package because the header deadline it exercises is a +// production constant measured in seconds, and a spec that waited it out would +// be the slowest in the suite. The seam is unexported for the same reason the +// worker tunnel's is: it is a test knob, not an operator knob. +var _ = Describe("A peer stream that never says what it wants", func() { + It("is refused rather than left holding a relay goroutine", func() { + // Without a deadline on the opening frame, a peer that opens a stream + // and then goes quiet parks a goroutine and a stream slot until the + // whole session dies. A peer need not be malicious to do it: a dialler + // killed between OpenStream and its first write leaves exactly this. + relay := newRelay(NewTunnelRegistry(nil, "me"), 50*time.Millisecond, 0) + store := NewSessionStore(relay.Stream) + DeferCleanup(store.CloseAll) + + a, b := net.Pipe() + accepted, err := yamux.Server(a, nil, nil) + Expect(err).ToNot(HaveOccurred()) + peer, err := yamux.Client(b, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { + _ = peer.Close() + _ = accepted.Close() + }) + store.Accept("peer-1", accepted) + + stream, err := peer.OpenStream(GinkgoT().Context()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + // Read with no deadline of our own: what is being asserted is that the + // RELAY answered, and a deadline here would be satisfied by a stream + // left parked just as well. + replies := make(chan error, 1) + go func() { + defer GinkgoRecover() + replies <- ReadRelayReply(stream) + }() + var reply error + Eventually(replies, "10s").Should(Receive(&reply)) + Expect(reply).To(MatchError(ErrRelayRequestInvalid)) + + ends := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := stream.Read(make([]byte, 1)) + ends <- err + }() + Eventually(ends, "10s").Should(Receive(HaveOccurred())) + }) +}) diff --git a/core/services/cluster/relay_test.go b/core/services/cluster/relay_test.go new file mode 100644 index 000000000000..389700b7d0e3 --- /dev/null +++ b/core/services/cluster/relay_test.go @@ -0,0 +1,291 @@ +// SPDX-License-Identifier: MIT + +package cluster_test + +import ( + "bytes" + "context" + "io" + "net" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// blockingRead runs one Read on its own goroutine with NO deadline set. +// +// The absence of the deadline is the point. A refusal and a stream left parked +// are indistinguishable to an assertion that waits for a deadline to expire: +// both produce an error at the same moment. Reading with no deadline at all +// means the channel only ever receives because the far side ANSWERED or ENDED +// the stream, so Eventually(...).Should(Receive()) is an assertion about the +// relay rather than about the clock. +func blockingRead(conn net.Conn) chan error { + done := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := conn.Read(make([]byte, 1)) + done <- err + }() + return done +} + +// relayReply reads the relay's answer, on its own goroutine and with no +// deadline, for the reason blockingRead gives. +func relayReply(conn net.Conn) chan error { + done := make(chan error, 1) + go func() { + defer GinkgoRecover() + done <- cluster.ReadRelayReply(conn) + }() + return done +} + +// readInto reads exactly len(buf) bytes on its own goroutine, with no deadline. +func readInto(conn net.Conn, buf []byte) chan error { + done := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := io.ReadFull(conn, buf) + done <- err + }() + return done +} + +// acceptOne hands back the next stream accepted on a session. +func acceptOne(sess *yamux.Session) chan net.Conn { + accepted := make(chan net.Conn, 1) + go func() { + defer GinkgoRecover() + stream, err := sess.AcceptStream() + if err != nil { + return + } + accepted <- stream + }() + return accepted +} + +var _ = Describe("The inter-replica relay", func() { + var ( + db *gorm.DB + reg *cluster.Registry + tun *cluster.TunnelRegistry + ctx context.Context + + // peer is the dialling replica's half of the peer link, the side a + // relayed request arrives from. + peer *yamux.Session + ) + + // openRelayStream opens a peer stream and names the node it is for. + openRelayStream := func(nodeID string) net.Conn { + GinkgoHelper() + stream, err := peer.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + Expect(cluster.WriteRelayRequest(stream, nodeID)).To(Succeed()) + return stream + } + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + tun = cluster.NewTunnelRegistry(reg, "me") + + store := cluster.NewSessionStore(cluster.NewRelay(tun).Stream) + DeferCleanup(store.CloseAll) + var accepted *yamux.Session + peer, accepted = yamuxPair() + store.Accept("peer-1", accepted) + }) + + It("splices a peer's stream onto a worker tunnel it holds, in both directions", func() { + // This is the whole point of the relay: with one tunnel per worker + // landing on ONE replica, every other replica reaches that worker only + // by relaying through this path, so with N replicas it carries roughly + // (N-1)/N of production traffic. + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + echoOnce(worker) + + stream := openRelayStream("w1") + Eventually(relayReply(stream), "10s").Should(Receive(BeNil())) + + _, err = stream.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + echoed := make([]byte, 4) + Eventually(readInto(stream, echoed), "10s").Should(Receive(BeNil())) + Expect(string(echoed)).To(Equal("ping")) + }) + + It("does not forward the frame it consumed, so the worker sees only the tunnelled protocol", func() { + // The relay request names the node for THIS hop and stops here. The + // worker's own request frame is written by the dialling replica and + // crosses untouched, so a relay that forwarded its own header would + // make every relayed stream unparseable at the worker while every + // locally-held one worked. + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + accepted := acceptOne(worker) + + stream := openRelayStream("w1") + Eventually(relayReply(stream), "10s").Should(Receive(BeNil())) + _, err = stream.Write([]byte("first")) + Expect(err).ToNot(HaveOccurred()) + + var workerSide net.Conn + Eventually(accepted, "10s").Should(Receive(&workerSide)) + first := make([]byte, 5) + Eventually(readInto(workerSide, first), "10s").Should(Receive(BeNil())) + Expect(string(first)).To(Equal("first")) + }) + + It("tears down the worker's side when the peer's side closes", func() { + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + accepted := acceptOne(worker) + + stream := openRelayStream("w1") + Eventually(relayReply(stream), "10s").Should(Receive(BeNil())) + var workerSide net.Conn + Eventually(accepted, "10s").Should(Receive(&workerSide)) + + Expect(stream.Close()).To(Succeed()) + // A relay that copies but does not tear down leaves a backend + // connection per abandoned request, and a worker runs out of them. + Eventually(blockingRead(workerSide), "10s").Should(Receive(HaveOccurred())) + }) + + It("tears down the peer's side when the worker's side closes", func() { + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + accepted := acceptOne(worker) + + stream := openRelayStream("w1") + Eventually(relayReply(stream), "10s").Should(Receive(BeNil())) + var workerSide net.Conn + Eventually(accepted, "10s").Should(Receive(&workerSide)) + + Expect(workerSide.Close()).To(Succeed()) + Eventually(blockingRead(stream), "10s").Should(Receive(MatchError(io.EOF))) + }) + + It("refuses a node it does not hold with the routing fact, and ENDS the stream", func() { + stream := openRelayStream("not-here") + + var reply error + Eventually(relayReply(stream), "10s").Should(Receive(&reply)) + Expect(reply).To(MatchError(cluster.ErrNotOwner)) + // Four conditions this phase forbids collapsing. ErrNotOwner says + // "ask the owner"; absence says "this worker is gone" and a scheduler + // acts on that; unreachability says "retry". + Expect(reply).ToNot(MatchError(cluster.ErrNoConnection)) + Expect(reply).ToNot(MatchError(cluster.ErrPeerUnreachable)) + Expect(reply).ToNot(MatchError(cluster.ErrInstanceNotFound)) + Expect(reply).ToNot(MatchError(cluster.ErrRelayUnavailable)) + + // Answering is not enough. A relay that says why and leaves the stream + // open has parked the peer on a request that will never be served, + // which reads as a slow replica rather than a refused request. + Eventually(blockingRead(stream), "10s").Should(Receive(MatchError(io.EOF))) + }) + + It("refuses rather than chasing a node another live replica owns", func() { + // A relay that resolved the owner and relayed onward would make a + // stale row into a loop between two replicas, each certain the other + // holds the worker. One hop, always: the dialling replica re-resolves. + Expect(reg.Register(ctx, "other", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "other") + Expect(err).ToNot(HaveOccurred()) + + stream := openRelayStream("w1") + var reply error + Eventually(relayReply(stream), "10s").Should(Receive(&reply)) + Expect(reply).To(MatchError(cluster.ErrNotOwner)) + Eventually(blockingRead(stream), "10s").Should(Receive(MatchError(io.EOF))) + }) + + It("reports a tunnel that will not carry a stream as infrastructure, never as not-owner", func() { + // The tunnel IS held here; its session died. Answering ErrNotOwner + // would send the dialling replica looking elsewhere for a worker that + // is attached right here, and it would find this replica again. + frontend, worker := workerTunnel() + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + Expect(worker.Close()).To(Succeed()) + Eventually(frontend.IsClosed, "10s").Should(BeTrue()) + + stream := openRelayStream("w1") + var reply error + Eventually(relayReply(stream), "10s").Should(Receive(&reply)) + Expect(reply).To(MatchError(cluster.ErrRelayUnavailable)) + Expect(reply).ToNot(MatchError(cluster.ErrNotOwner)) + Expect(reply).ToNot(MatchError(cluster.ErrNoConnection)) + Expect(reply).ToNot(MatchError(cluster.ErrInstanceNotFound)) + Eventually(blockingRead(stream), "10s").Should(Receive(MatchError(io.EOF))) + }) + + It("refuses a malformed opening frame as the caller's bug", func() { + stream, err := peer.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + // A well-formed frame carrying no node id. A relay that read this as a + // node named "" would go looking for it and answer ErrNotOwner, which + // tells the caller to retry elsewhere for a request no replica can + // ever serve. + _, err = stream.Write([]byte{0x00, 0x00}) + Expect(err).ToNot(HaveOccurred()) + + var reply error + Eventually(relayReply(stream), "10s").Should(Receive(&reply)) + Expect(reply).To(MatchError(cluster.ErrRelayRequestInvalid)) + Expect(reply).ToNot(MatchError(cluster.ErrNotOwner)) + Expect(reply).ToNot(MatchError(cluster.ErrRelayUnavailable)) + Eventually(blockingRead(stream), "10s").Should(Receive(MatchError(io.EOF))) + }) +}) + +var _ = Describe("The relay wire framing", func() { + // The relay hop and the worker tunnel hop travel back to back on one + // stream, and a dialler reads a reply from each in order. Giving them + // disjoint vocabularies means a reader applied to the wrong hop fails + // loudly rather than returning a plausible sentinel for the other hop, + // which would report a worker's refusal as the owning replica's and send a + // retry to the wrong place. + It("does not read a relay reply as a worker tunnel reply", func() { + frame := &bytes.Buffer{} + Expect(cluster.WriteRelayAccepted(frame)).To(Succeed()) + Expect(cluster.ReadStreamReply(frame)).To(HaveOccurred()) + }) + + It("does not read a worker tunnel reply as a relay reply", func() { + frame := &bytes.Buffer{} + Expect(cluster.WriteStreamAccepted(frame)).To(Succeed()) + Expect(cluster.ReadRelayReply(frame)).To(HaveOccurred()) + }) + + It("round-trips a node id", func() { + frame := &bytes.Buffer{} + Expect(cluster.WriteRelayRequest(frame, "node-7")).To(Succeed()) + nodeID, err := cluster.ReadRelayRequest(frame) + Expect(err).ToNot(HaveOccurred()) + Expect(nodeID).To(Equal("node-7")) + }) + + It("refuses to write an empty node id, rather than spending a round trip on it", func() { + Expect(cluster.WriteRelayRequest(&bytes.Buffer{}, "")).To(HaveOccurred()) + }) +}) diff --git a/core/services/cluster/sessions.go b/core/services/cluster/sessions.go index 7ee5fa3327af..afcdd1b97384 100644 --- a/core/services/cluster/sessions.go +++ b/core/services/cluster/sessions.go @@ -23,6 +23,10 @@ type SessionStore struct { // onStream handles one accepted stream and owns closing it. A nil handler // closes the stream immediately, which is what a replica with no relay // installed should do: refuse promptly rather than leave a peer parked. + // + // In distributed mode this is Relay.Stream, which splices the stream onto + // a worker tunnel this replica holds. Nil is reached only from specs, and + // from a caller that wants a store with no relay. onStream func(peerID string, stream net.Conn) mu sync.Mutex diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go index 604bbfec93d4..691586c254ca 100644 --- a/core/services/cluster/splice.go +++ b/core/services/cluster/splice.go @@ -80,29 +80,27 @@ func copyStream(dst io.Writer, src io.Reader) error { // *net.TCPConn takes a WriteTo/ReadFrom path that would hand one back. So a // bare io.EOF arriving here came from a failing Write or Close, where it means // the peer is gone, and yamux produces exactly that when a Write races its -// session's shutdown (see isMuxSessionFailure). +// session's shutdown (see isMuxFailure). // // A socket-level abort (ECONNRESET, EPIPE) is deliberately absent too, which // makes the same underlying event, a peer aborting mid-stream, reach the caller -// as nil over a yamux tunnel and as an error over a raw socket. That asymmetry -// is intended: the yamux endings are the teardown Splice's own Close provokes, -// so this primitive is the only thing that can tell them from a fault, whereas -// whether an aborted request is routine or a failure is the relay's policy and -// only the relay knows which request was abandoned. +// as nil over a raw socket where it would be an error. That asymmetry is +// narrower than it was, since a yamux abort is now reported (see isMuxFailure), +// and what remains of it is that a raw socket cannot say who aborted. // // The mux checks run first, and that ordering is load-bearing: a dying yamux // session usually hands every live stream its own cause wrapped up -// (session.go:330), and that cause is routinely a closed-socket error, so -// consulting the generic endings first would report a peer that vanished -// mid-request as a clean completion. +// (go-yamux/v5@v5.1.0/session.go:328-337), and that cause is routinely a +// closed-socket error, so consulting the generic endings first would report a +// peer that vanished mid-request as a clean completion. func normalizeStreamErr(err error) error { if err == nil { return nil } - if isMuxSessionFailure(err) { + if isMuxFailure(err) { return err } - if isMuxStreamTeardown(err) { + if isMuxLocalTeardown(err) { return nil } if errors.Is(err, net.ErrClosed) || @@ -116,23 +114,66 @@ func normalizeStreamErr(err error) error { // declared with it because the constant itself is unexported. var normalGoAwayCode = yamux.ErrRemoteGoAway.ErrorCode -// isMuxSessionFailure reports whether err is the yamux session underneath a -// stream dying, as opposed to a single stream being torn down. The distinction -// matters because Splice must stay quiet about the teardown it provokes itself -// while still reporting a dead peer: a keepalive timeout, a broken TCP -// connection or a protocol error under a relayed request has to reach the -// caller, or a failed inference looks like a finished one. -func isMuxSessionFailure(err error) bool { - // A go-away ends the whole session. Only the "no error" code is a normal - // ending; a protocol or internal error go-away is a real failure. +// isMuxFailure reports whether err is yamux saying the conversation was CUT, +// as opposed to ended by this side. +// +// The distinction is what keeps Splice quiet about the teardown it provokes +// itself while still reporting a request that died: a keepalive timeout, a +// broken connection, a peer that reset the stream or a peer that went away +// under a relayed request has to reach the caller, or a failed inference looks +// like a finished one. +// +// TWO OF THESE ARE THE POLICY PHASE 1 LEFT OPEN, and this is where they are +// settled, by the relay in core/services/cluster/relay.go, which is Splice's +// first production caller. Both used to be reported as normal termination. +// Neither can be settled by a caller reading Splice's result, because a result +// mapped to nil carries nothing left to reclassify, so the decision has to live +// at the classifier; the two callers there are the relay and the worker tunnel, +// and both are splicing an in-flight request, so both want the same answer. +// +// 1. A peer-initiated stream reset, *StreamError{Remote: true}. yamux builds +// it in processFlags when an RST frame arrives on the stream +// (stream.go:432-449); a reset this side asked for carries Remote: false +// instead (stream.go:283-291), and Splice never resets anything anyway, its +// own Close sending a FIN (stream.go:303-331, 365-368). So Remote: true is +// unambiguously "the far side aborted this stream", which for a relayed +// request means the response was truncated. REPORTED. The caller decides +// how loud that is: a client cancelling produces one per cancellation, so +// the relay logs it at debug rather than treating it as a fault. +// +// 2. A graceful go-away from the peer, ErrRemoteGoAway. handleGoAway returns +// it for code goAwayNormal (session.go:829-833), recv closes the session +// with it, and close hands it UNWRAPPED to every live stream, because it +// already is a *GoAwayError and so escapes the ErrStreamReset wrapping +// (session.go:328-337, stream.go:371-387). Graceful describes the SESSION, +// not the requests on it: every one of those streams was mid-request. +// REPORTED, for the same reason as above. +// +// The locally-initiated forms of both stay silent, and keying on Remote is what +// separates them: ErrSessionShutdown is a *GoAwayError with Remote: false +// (const.go:96) and is exactly what this process closing its own session +// produces (session.go:284). +// +// What this does NOT do is make the far side see a failure. Splice ends both +// streams with Close, which is a FIN, and a reset after that is a no-op because +// Close has already moved the stream to streamFinished (stream.go:266-272, +// 303-331, 336-361). Propagating a truncation as an RST would mean reshaping +// Splice's teardown, and it buys little: the protocols relayed here are gRPC +// and HTTP, both of which detect a body that ended without its trailers or its +// final chunk. Reporting is what the caller needs and this is where it comes +// from. +func isMuxFailure(err error) bool { + // A go-away ends the whole session. Only a normal-code go-away this side + // sent is a normal ending. var goAway *yamux.GoAwayError if errors.As(err, &goAway) { - return goAway.ErrorCode != normalGoAwayCode + return goAway.Remote || goAway.ErrorCode != normalGoAwayCode } - // A stream error is scoped to one stream, whatever killed it. + // A stream error is scoped to one stream. Only a reset this side asked for + // is a normal ending. var streamErr *yamux.StreamError if errors.As(err, &streamErr) { - return false + return streamErr.Remote } // Session.close gives every stream it kills ErrStreamReset wrapped around // the cause, so the bare sentinel means this stream was reset and a @@ -142,32 +183,37 @@ func isMuxSessionFailure(err error) bool { // Wrapped is not the only way a dead session shows up, though. close() // publishes shutdownErr and closes shutdownCh before it force-closes the // streams, so a Write or Close landing in that window gets the raw cause - // back instead (session.go:507-510, 528-533). That form is unrecognisable + // back instead (session.go:305-308, 528-533). That form is unrecognisable // as yamux at all, which is why normalizeStreamErr no longer forgives a // bare io.EOF: for a peer that vanished, the raw cause is precisely io.EOF. return errors.Is(err, yamux.ErrStreamReset) && err != yamux.ErrStreamReset } -// isMuxStreamTeardown reports whether err is yamux ending one stream, which -// Splice mostly provokes itself: closing a stream whose session has already -// shut down normally returns ErrSessionShutdown from the FIN write, and a copy -// parked on a stream that gets closed comes back with ErrStreamClosed or a -// reset. A reset does not only mean that, though; the same sentinel heads the -// error a dying session hands its streams, which is why isMuxSessionFailure -// runs first. None of yamux's error types match net.ErrClosed, so all of this -// has to be recognised by shape. -func isMuxStreamTeardown(err error) bool { +// isMuxLocalTeardown reports whether err is yamux ending one stream at this +// side's request, which Splice mostly provokes itself: closing a stream whose +// session has already shut down normally returns ErrSessionShutdown from the +// FIN write, and a copy parked on a stream that gets closed comes back with +// ErrStreamClosed from a Write (stream.go:157-159) or the bare ErrStreamReset +// from a Read, which is what CloseRead installs (stream.go:348-349). A reset +// does not only mean that, though; the same sentinel heads the error a dying +// session hands its streams, and the remote forms belong to isMuxFailure, which +// is why that runs first. None of yamux's error types match net.ErrClosed, so +// all of this has to be recognised by shape. +func isMuxLocalTeardown(err error) bool { // Sentinels by identity, never errors.Is: the wrapped forms belong to a // dead session and are reported instead. ErrSessionShutdown is absent on // purpose rather than by oversight, being a *GoAwayError carrying the - // normal code, which the last check below covers. + // normal code with Remote false, which the last check below covers. if err == yamux.ErrStreamClosed || err == yamux.ErrStreamReset { return true } + // Remote is re-checked rather than assumed from the ordering, so that this + // predicate is true to its own name if it is ever called from anywhere + // else. var streamErr *yamux.StreamError if errors.As(err, &streamErr) { - return true + return !streamErr.Remote } var goAway *yamux.GoAwayError - return errors.As(err, &goAway) && goAway.ErrorCode == normalGoAwayCode + return errors.As(err, &goAway) && !goAway.Remote && goAway.ErrorCode == normalGoAwayCode } diff --git a/core/services/cluster/splice_test.go b/core/services/cluster/splice_test.go index 7b423a142730..7b1001cd8601 100644 --- a/core/services/cluster/splice_test.go +++ b/core/services/cluster/splice_test.go @@ -172,8 +172,12 @@ var _ = Describe("Splice", func() { Entry("a closed yamux stream", yamux.ErrStreamClosed), Entry("a reset yamux stream", yamux.ErrStreamReset), Entry("a shut-down yamux session", yamux.ErrSessionShutdown), - Entry("a stream reset by the remote", &yamux.StreamError{ErrorCode: 1, Remote: true}), - Entry("a go-away from the remote", yamux.ErrRemoteGoAway), + // The LOCAL forms of the two endings the relay settled below. They + // stay normal because they are the teardown this side asked for, + // and Remote is the only thing separating them from the endings a + // peer inflicts. + Entry("a stream this side reset", &yamux.StreamError{ErrorCode: 0, Remote: false}), + Entry("a go-away this side sent", &yamux.GoAwayError{ErrorCode: 0, Remote: false}), ) // sessionDeath is the exact shape Session.close hands every live @@ -205,6 +209,14 @@ var _ = Describe("Splice", func() { Entry("a peer that vanished", sessionDeath(io.EOF)), Entry("a protocol-error go-away", &yamux.GoAwayError{Remote: true, ErrorCode: 1}), Entry("an internal-error go-away", &yamux.GoAwayError{Remote: true, ErrorCode: 2}), + // The two endings phase 1 left open and the relay, Splice's first + // production caller, settled as failures. Both truncate whatever + // was in flight, and reporting them as normal termination is how a + // half-finished inference comes to look like a short one that + // completed. See isMuxFailure for why the decision could not be + // left to a caller reading Splice's result. + Entry("a stream the peer reset", &yamux.StreamError{ErrorCode: 1, Remote: true}), + Entry("a graceful go-away from the peer", yamux.ErrRemoteGoAway), ) // A bare io.EOF can only reach Splice from a failing Write. io.Copy From 1036f5664e0279614ae41a143d5c23481d0ffdb2 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 14:39:24 +0000 Subject: [PATCH 30/42] fix(cluster): express the splice policy once, and pin the relay's budgets Review round 1 on task 5. Eight non-blocking items, all addressed. The classifier read Remote in two predicates with a report-by-default fallthrough behind them, so reverting either read left the whole suite green: the error reached the same answer down the other path. A correctness argument that rests on mutation evidence cannot afford a shape that cannot be mutated in pieces, so the two predicates collapse into one muxVerdict deciding each error type once. Falsifying either Remote read now reddens exactly one spec. Three claims the comments made loudly and nothing tested: - clearing the header read deadline before the splice. Deleting the clear left all 49 focused specs green, while in production it is the difference between a relayed response that streams for an hour and one that dies after fifteen seconds of quiet; - the open budget bounding the open and nothing after it; - closing the worker-side stream when the acceptance reply cannot be written, which leaks one stream on the worker per failure. All three are pinned now. The first two share a spec that sets both budgets to 50ms and then watches the conversation outlive them by ten times, which is an assertion about an event that must not happen and so is the one wait a channel cannot replace. The third drives the relay with a peer stream that delivers a request and then fails every write, because no pair of live yamux sessions can be made to fail that write on cue. The disjoint-vocabulary argument was specced for the accepted frame only. Both refusal directions are covered now, and asserted as "not one of the other hop's sentinels" rather than merely "an error", since reading a relay refusal with the tunnel's reader always errors and the question is whether it errors as the wrong thing. The open budget stays non-configurable, and says so: the number that matters is how long the original client will wait, which is not known on this side and is not something a deployment-wide constant can stand in for. The honest fix is the caller's remaining budget travelling in the request frame, which belongs to the dialler that has the budget. Two comments corrected: nothing deadlines the peer stream after the clear, so the tunnelled protocol's own deadlines cannot be what justifies clearing it; and the membership sweep deletes departed replicas but reports only how many, so identifying them is work that would have to be done, not knowledge waiting to be plumbed. Recorded at muxVerdict: a remote RST that does not ride a typeWindowUpdate frame yields the bare sentinel and is still silenced, which is unreachable between two go-yamux peers but keeps the new rule from reading as unconditional. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Opus 5 [claude-code] --- core/services/cluster/peerlink.go | 8 +- core/services/cluster/relay.go | 21 +- core/services/cluster/relay_internal_test.go | 245 +++++++++++++++++++ core/services/cluster/relay_test.go | 35 +++ core/services/cluster/splice.go | 112 +++++---- 5 files changed, 359 insertions(+), 62 deletions(-) diff --git a/core/services/cluster/peerlink.go b/core/services/cluster/peerlink.go index 047fcbc5472a..5f5e3186445b 100644 --- a/core/services/cluster/peerlink.go +++ b/core/services/cluster/peerlink.go @@ -233,10 +233,10 @@ func (p *PeerPool) Open(ctx context.Context, peerID string) (net.Conn, error) { // has left the deployment but is still listening keeps a live WebSocket and the // two yamux loop goroutines behind it for as long as this process runs; a peer // that is genuinely gone is reclaimed by the 30s keepalive default, so the real -// exposure is narrow. There is no Forget: the membership sweep does know which -// replicas have left, but nothing plumbs that knowledge to this pool, so an -// entry for a departed peer outlives it and only the keepalive reclaims what it -// holds. +// exposure is narrow. There is no Forget: the membership sweep DELETES departed +// replicas but reports only how many, so which ones they were would have to be +// surfaced before anything could be plumbed here. Until it is, an entry for a +// departed peer outlives it and only the keepalive reclaims what it holds. func (p *PeerPool) link(peerID string) (*peerLink, error) { p.mu.Lock() defer p.mu.Unlock() diff --git a/core/services/cluster/relay.go b/core/services/cluster/relay.go index 99164c6fbd74..9fa03807d465 100644 --- a/core/services/cluster/relay.go +++ b/core/services/cluster/relay.go @@ -178,6 +178,15 @@ const ( // which is the only reason there is one here. Without the bound, a worker // that has stopped accepting would turn a refusable condition into a parked // peer, which is the one outcome this path exists to avoid. + // + // It is deliberately NOT configurable. No operator has the information to + // set it: the number that matters is how long the ORIGINAL client is + // willing to wait, which is not known on this side of the link and is not + // something a deployment-wide constant can stand in for. The honest fix is + // the caller's remaining budget travelling in the relay request frame, and + // that belongs to the dialler that has the budget. Until then this is a + // backstop against parking, generous on purpose, because refusing healthy + // traffic costs more than waiting. relayOpenTimeout = 15 * time.Second ) @@ -271,10 +280,14 @@ func (r *Relay) accept(peerID string, stream net.Conn) (net.Conn, bool) { return nil, false } - // Cleared before the open rather than after the reply: everything past this - // frame belongs to the worker tunnel's conversation, which brings its own - // deadlines, and one left armed here would abort a long inference stream in - // the middle. + // Cleared before the open rather than after the reply, and NOTHING arms + // another deadline on this stream afterwards. That is the intent rather + // than an omission: what follows is a relayed request whose length is the + // caller's business, and a header deadline left armed here would abort a + // long inference stream after any quiet moment in the middle of it. What + // still bounds the conversation is the peer link's own keepalive, which + // kills the session under it when the far side stops answering, and + // whatever deadline the original client is holding. 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 diff --git a/core/services/cluster/relay_internal_test.go b/core/services/cluster/relay_internal_test.go index 855ac58653bf..e23f6c55c888 100644 --- a/core/services/cluster/relay_internal_test.go +++ b/core/services/cluster/relay_internal_test.go @@ -3,9 +3,16 @@ package cluster import ( + "bytes" + "context" + "errors" + "io" "net" + "sync" "time" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/libp2p/go-yamux/v5" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -61,3 +68,241 @@ var _ = Describe("A peer stream that never says what it wants", func() { Eventually(ends, "10s").Should(Receive(HaveOccurred())) }) }) + +// backloggedPair returns a peer/relay session pair whose SYN backlog is one +// stream deep, so a single un-accepted open fills it and the next one parks. +// yamux's default is 256 (mux.go, DefaultConfig), and filling that from a spec +// would mean 256 real opens to prove one property. +func backloggedPair(backlog int) (dialled, accepted *yamux.Session) { + GinkgoHelper() + cfg := yamux.DefaultConfig() + cfg.AcceptBacklog = backlog + a, b := net.Pipe() + var err error + accepted, err = yamux.Server(a, cfg, nil) + Expect(err).ToNot(HaveOccurred()) + dialled, err = yamux.Client(b, cfg, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { + _ = dialled.Close() + _ = accepted.Close() + }) + return dialled, accepted +} + +// unwritableStream is a peer stream that delivers one relay request and then +// fails every write. It stands in for a peer that vanished between opening the +// stream and hearing the answer, which is the only way the acceptance reply +// fails, and which no pair of live yamux sessions can be made to do on cue. +type unwritableStream struct { + net.Conn + request []byte + read int + closed chan struct{} + closeOne sync.Once +} + +func newUnwritableStream(nodeID string) *unwritableStream { + GinkgoHelper() + frame := &bytes.Buffer{} + Expect(WriteRelayRequest(frame, nodeID)).To(Succeed()) + return &unwritableStream{request: frame.Bytes(), closed: make(chan struct{})} +} + +func (s *unwritableStream) Read(p []byte) (int, error) { + if s.read >= len(s.request) { + // Never EOF: an EOF here would end the relay for a reason other than + // the failed write, and the spec would pass without exercising it. + <-s.closed + return 0, io.EOF + } + n := copy(p, s.request[s.read:]) + s.read += n + return n, nil +} + +func (s *unwritableStream) Write([]byte) (int, error) { return 0, errors.New("peer went away") } + +func (s *unwritableStream) Close() error { + s.closeOne.Do(func() { close(s.closed) }) + return nil +} + +func (s *unwritableStream) SetReadDeadline(time.Time) error { return nil } + +var _ = Describe("The relay's own budgets", func() { + var ( + reg *Registry + tun *TunnelRegistry + ctx context.Context + ) + + BeforeEach(func() { + ctx = context.Background() + db := testutil.SetupTestDB() + Expect(Migrate(ctx, db)).To(Succeed()) + reg = NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + tun = NewTunnelRegistry(reg, "me") + }) + + It("stops bounding the stream once the relay hands it over", func() { + // Both budgets are set to 50ms here and both are deliberately shorter + // than the window this spec then watches. A header deadline left armed + // past acceptance, or an open budget applied to the stream it produced, + // would abort a relayed inference after 50ms of quiet, which in + // production is the difference between a response that streams for an + // hour and one that dies mid-token. + relay := newRelay(tun, 50*time.Millisecond, 50*time.Millisecond) + store := NewSessionStore(relay.Stream) + DeferCleanup(store.CloseAll) + peer, accepted := backloggedPair(256) + store.Accept("peer-1", accepted) + + worker, frontend := backloggedPair(256) + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + workerSide := make(chan net.Conn, 1) + go func() { + defer GinkgoRecover() + stream, err := worker.AcceptStream() + if err != nil { + return + } + workerSide <- stream + }() + + stream, err := peer.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + Expect(WriteRelayRequest(stream, "w1")).To(Succeed()) + + replies := make(chan error, 1) + go func() { + defer GinkgoRecover() + replies <- ReadRelayReply(stream) + }() + Eventually(replies, "10s").Should(Receive(BeNil())) + + var served net.Conn + Eventually(workerSide, "10s").Should(Receive(&served)) + + // One reader for both questions, so that watching for a teardown does + // not eat the bytes the second half of the spec is waiting for. + data := make(chan []byte, 4) + ended := make(chan error, 1) + go func() { + defer GinkgoRecover() + buf := make([]byte, 64) + for { + n, err := served.Read(buf) + if n > 0 { + data <- append([]byte(nil), buf[:n]...) + } + if err != nil { + ended <- err + return + } + } + }() + + // An assertion about an event that must NOT happen, which is the one + // kind a channel cannot replace: a torn-down splice ends the worker's + // side, and there is no event for "still alive". The window is ten + // times the budgets it is watching. + Consistently(ended, "500ms", "50ms").ShouldNot(Receive(), + "the relay tore the stream down on a budget that should have stopped applying at acceptance") + + // And it is not merely un-torn-down: it still carries bytes, long after + // both budgets would have expired. + _, err = stream.Write([]byte("late")) + Expect(err).ToNot(HaveOccurred()) + Eventually(data, "10s").Should(Receive(Equal([]byte("late")))) + }) + + It("refuses rather than parking when the worker's tunnel will not take another stream", func() { + // The open budget exists because yamux BLOCKS an open once the accept + // backlog is full rather than failing it, so without a bound an + // overloaded worker turns a refusable condition into a parked peer. + relay := newRelay(tun, 0, 50*time.Millisecond) + store := NewSessionStore(relay.Stream) + DeferCleanup(store.CloseAll) + peer, accepted := backloggedPair(256) + store.Accept("peer-1", accepted) + + _, frontend := backloggedPair(1) + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + // One un-accepted open fills the one-deep backlog; the relay's own open + // is the one that has to wait. + filler, err := frontend.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = filler.Close() }) + + stream, err := peer.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + Expect(WriteRelayRequest(stream, "w1")).To(Succeed()) + + replies := make(chan error, 1) + go func() { + defer GinkgoRecover() + replies <- ReadRelayReply(stream) + }() + var reply error + Eventually(replies, "10s").Should(Receive(&reply)) + Expect(reply).To(MatchError(ErrRelayUnavailable)) + // The tunnel IS held here, so this must not read as a routing fact. + Expect(reply).ToNot(MatchError(ErrNotOwner)) + + ends := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := stream.Read(make([]byte, 1)) + ends <- err + }() + Eventually(ends, "10s").Should(Receive(HaveOccurred())) + }) + + It("closes the worker's stream when it cannot tell the peer the stream was accepted", func() { + // The reply is the last thing that can fail after a worker stream has + // been opened. A relay that gave up without closing it would leak one + // stream on the worker per failed acceptance, and the worker cannot + // tell those from live ones. + relay := newRelay(tun, 0, 0) + worker, frontend := backloggedPair(256) + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + + workerSide := make(chan net.Conn, 1) + go func() { + defer GinkgoRecover() + stream, err := worker.AcceptStream() + if err != nil { + return + } + workerSide <- stream + }() + + peerStream := newUnwritableStream("w1") + done := make(chan struct{}) + go func() { + defer GinkgoRecover() + defer close(done) + relay.Stream("peer-1", peerStream) + }() + Eventually(done, "10s").Should(BeClosed()) + + var served net.Conn + Eventually(workerSide, "10s").Should(Receive(&served)) + ended := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := served.Read(make([]byte, 1)) + ended <- err + }() + Eventually(ended, "10s").Should(Receive(HaveOccurred()), + "the worker-side stream outlived the relay that opened it") + }) +}) diff --git a/core/services/cluster/relay_test.go b/core/services/cluster/relay_test.go index 389700b7d0e3..e0c7bd0499c2 100644 --- a/core/services/cluster/relay_test.go +++ b/core/services/cluster/relay_test.go @@ -288,4 +288,39 @@ var _ = Describe("The relay wire framing", func() { It("refuses to write an empty node id, rather than spending a round trip on it", func() { Expect(cluster.WriteRelayRequest(&bytes.Buffer{}, "")).To(HaveOccurred()) }) + + // Acceptance is not the whole surface. A refusal read by the wrong hop's + // reader must not come back as one of that hop's own sentinels: "the + // owning replica does not hold this worker" arriving as "the worker does + // not serve that tag" would send a retry to the wrong end of the path, and + // it would look like a perfectly ordinary answer on the way. + DescribeTable("does not read a relay refusal as one of the worker tunnel's", + func(reason error) { + frame := &bytes.Buffer{} + Expect(cluster.WriteRelayRefusal(frame, reason)).To(Succeed()) + err := cluster.ReadStreamReply(frame) + Expect(err).To(HaveOccurred()) + Expect(err).ToNot(MatchError(cluster.ErrStreamTagUnknown)) + Expect(err).ToNot(MatchError(cluster.ErrStreamTargetUnavailable)) + Expect(err).ToNot(MatchError(cluster.ErrStreamRequestInvalid)) + }, + Entry("not the owner", cluster.ErrNotOwner), + Entry("the tunnel will not carry a stream", cluster.ErrRelayUnavailable), + Entry("a malformed relay request", cluster.ErrRelayRequestInvalid), + ) + + DescribeTable("does not read a worker tunnel refusal as one of the relay's", + func(reason error) { + frame := &bytes.Buffer{} + Expect(cluster.WriteStreamRefusal(frame, reason)).To(Succeed()) + err := cluster.ReadRelayReply(frame) + Expect(err).To(HaveOccurred()) + Expect(err).ToNot(MatchError(cluster.ErrNotOwner)) + Expect(err).ToNot(MatchError(cluster.ErrRelayUnavailable)) + Expect(err).ToNot(MatchError(cluster.ErrRelayRequestInvalid)) + }, + Entry("an unknown stream tag", cluster.ErrStreamTagUnknown), + Entry("a local service that will not answer", cluster.ErrStreamTargetUnavailable), + Entry("a malformed stream request", cluster.ErrStreamRequestInvalid), + ) }) diff --git a/core/services/cluster/splice.go b/core/services/cluster/splice.go index 691586c254ca..32fe4358f3d7 100644 --- a/core/services/cluster/splice.go +++ b/core/services/cluster/splice.go @@ -80,27 +80,27 @@ func copyStream(dst io.Writer, src io.Reader) error { // *net.TCPConn takes a WriteTo/ReadFrom path that would hand one back. So a // bare io.EOF arriving here came from a failing Write or Close, where it means // the peer is gone, and yamux produces exactly that when a Write races its -// session's shutdown (see isMuxFailure). +// session's shutdown (see muxVerdict). // // A socket-level abort (ECONNRESET, EPIPE) is deliberately absent too, which // makes the same underlying event, a peer aborting mid-stream, reach the caller // as nil over a raw socket where it would be an error. That asymmetry is -// narrower than it was, since a yamux abort is now reported (see isMuxFailure), +// narrower than it was, since a yamux abort is now reported (see muxVerdict), // and what remains of it is that a raw socket cannot say who aborted. // -// The mux checks run first, and that ordering is load-bearing: a dying yamux -// session usually hands every live stream its own cause wrapped up -// (go-yamux/v5@v5.1.0/session.go:328-337), and that cause is routinely a +// The mux verdict is consulted first, and that ordering is load-bearing: a +// dying yamux session usually hands every live stream its own cause wrapped up +// (go-yamux/v5@v5.1.0 session.go, Session.close), and that cause is routinely a // closed-socket error, so consulting the generic endings first would report a // peer that vanished mid-request as a clean completion. func normalizeStreamErr(err error) error { if err == nil { return nil } - if isMuxFailure(err) { - return err - } - if isMuxLocalTeardown(err) { + if recognised, report := muxVerdict(err); recognised { + if report { + return err + } return nil } if errors.Is(err, net.ErrClosed) || @@ -114,14 +114,24 @@ func normalizeStreamErr(err error) error { // declared with it because the constant itself is unexported. var normalGoAwayCode = yamux.ErrRemoteGoAway.ErrorCode -// isMuxFailure reports whether err is yamux saying the conversation was CUT, -// as opposed to ended by this side. +// muxVerdict classifies a yamux ending. recognised says the error came from the +// multiplexer at all; report says the ending was INFLICTED on this stream +// rather than asked for by this side. // -// The distinction is what keeps Splice quiet about the teardown it provokes -// itself while still reporting a request that died: a keepalive timeout, a -// broken connection, a peer that reset the stream or a peer that went away -// under a relayed request has to reach the caller, or a failed inference looks -// like a finished one. +// It is ONE function, and that is the point rather than a matter of taste. The +// policy below turns on a single bit, Remote, and an earlier shape read that +// bit in two predicates with a report-by-default fallthrough behind them. +// Reverting either read left the whole suite green, because the error reached +// the same answer down the other path: the classifier could not be +// mutation-tested in pieces, which in code whose correctness argument IS its +// mutation evidence is worse than the duplication it bought. Here each type is +// decided once, so falsifying either read reddens a spec. +// +// The distinction it draws is what keeps Splice quiet about the teardown it +// provokes itself while still reporting a request that died: a keepalive +// timeout, a broken connection, a peer that reset the stream or a peer that +// went away under a relayed request has to reach the caller, or a failed +// inference looks like a finished one. // // TWO OF THESE ARE THE POLICY PHASE 1 LEFT OPEN, and this is where they are // settled, by the relay in core/services/cluster/relay.go, which is Splice's @@ -154,6 +164,14 @@ var normalGoAwayCode = yamux.ErrRemoteGoAway.ErrorCode // (const.go:96) and is exactly what this process closing its own session // produces (session.go:284). // +// The rule is "a remote reset is reported" and not "every remote reset is +// reported". yamux only builds a *StreamError when the RST rides a +// typeWindowUpdate frame (stream.go:436); an RST on any other frame type +// yields the BARE ErrStreamReset sentinel, which is claimed below as this +// side's own teardown and silenced. Every reset go-yamux itself sends uses +// typeWindowUpdate, so the gap is unreachable between two LocalAI processes and +// only a foreign multiplexer implementation could reach it. +// // What this does NOT do is make the far side see a failure. Splice ends both // streams with Close, which is a FIN, and a reset after that is a no-op because // Close has already moved the stream to streamFinished (stream.go:266-272, @@ -162,58 +180,44 @@ var normalGoAwayCode = yamux.ErrRemoteGoAway.ErrorCode // and HTTP, both of which detect a body that ended without its trailers or its // final chunk. Reporting is what the caller needs and this is where it comes // from. -func isMuxFailure(err error) bool { +func muxVerdict(err error) (recognised, report bool) { // A go-away ends the whole session. Only a normal-code go-away this side // sent is a normal ending. var goAway *yamux.GoAwayError if errors.As(err, &goAway) { - return goAway.Remote || goAway.ErrorCode != normalGoAwayCode + return true, goAway.Remote || goAway.ErrorCode != normalGoAwayCode } // A stream error is scoped to one stream. Only a reset this side asked for // is a normal ending. var streamErr *yamux.StreamError if errors.As(err, &streamErr) { - return streamErr.Remote + return true, streamErr.Remote } - // Session.close gives every stream it kills ErrStreamReset wrapped around - // the cause, so the bare sentinel means this stream was reset and a - // wrapped one means the session died under it. Identity is what separates - // them; errors.Is cannot. + // Sentinels by identity, never errors.Is, and before the wrapped check + // below: these are the endings Splice provokes itself. Closing a stream + // whose session has already shut down normally returns ErrSessionShutdown + // from the FIN write, which the go-away branch above has already claimed; + // a copy parked on a stream that gets closed comes back with + // ErrStreamClosed from a Write (stream.go:157-159) or the bare + // ErrStreamReset from a Read, which is what CloseRead installs + // (stream.go:348-349). + if err == yamux.ErrStreamClosed || err == yamux.ErrStreamReset { + return true, false + } + // The same sentinel WRAPPED means something else entirely: Session.close + // gives every stream it kills ErrStreamReset wrapped around the cause, so + // this is the session dying under a live stream. Identity above is what + // separates the two; errors.Is cannot. // // Wrapped is not the only way a dead session shows up, though. close() // publishes shutdownErr and closes shutdownCh before it force-closes the // streams, so a Write or Close landing in that window gets the raw cause // back instead (session.go:305-308, 528-533). That form is unrecognisable - // as yamux at all, which is why normalizeStreamErr no longer forgives a - // bare io.EOF: for a peer that vanished, the raw cause is precisely io.EOF. - return errors.Is(err, yamux.ErrStreamReset) && err != yamux.ErrStreamReset -} - -// isMuxLocalTeardown reports whether err is yamux ending one stream at this -// side's request, which Splice mostly provokes itself: closing a stream whose -// session has already shut down normally returns ErrSessionShutdown from the -// FIN write, and a copy parked on a stream that gets closed comes back with -// ErrStreamClosed from a Write (stream.go:157-159) or the bare ErrStreamReset -// from a Read, which is what CloseRead installs (stream.go:348-349). A reset -// does not only mean that, though; the same sentinel heads the error a dying -// session hands its streams, and the remote forms belong to isMuxFailure, which -// is why that runs first. None of yamux's error types match net.ErrClosed, so -// all of this has to be recognised by shape. -func isMuxLocalTeardown(err error) bool { - // Sentinels by identity, never errors.Is: the wrapped forms belong to a - // dead session and are reported instead. ErrSessionShutdown is absent on - // purpose rather than by oversight, being a *GoAwayError carrying the - // normal code with Remote false, which the last check below covers. - if err == yamux.ErrStreamClosed || err == yamux.ErrStreamReset { - return true - } - // Remote is re-checked rather than assumed from the ordering, so that this - // predicate is true to its own name if it is ever called from anywhere - // else. - var streamErr *yamux.StreamError - if errors.As(err, &streamErr) { - return !streamErr.Remote + // as yamux at all, and is why it is left unrecognised here rather than + // guessed at: normalizeStreamErr no longer forgives a bare io.EOF, because + // for a peer that vanished the raw cause is precisely io.EOF. + if errors.Is(err, yamux.ErrStreamReset) { + return true, true } - var goAway *yamux.GoAwayError - return errors.As(err, &goAway) && !goAway.Remote && goAway.ErrorCode == normalGoAwayCode + return false, false } From 75953d9f632baf81004f91a3a0e720005ec6b3bf Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 15:27:10 +0000 Subject: [PATCH 31/42] feat(cluster): reach every worker through its tunnel, never its address The tunnel, the fence, the registry and the relay were all built and none of them carried a byte: every dial from the frontend still went to the address a worker registered. This is where that stops. One WorkerDialer resolves where a worker's tunnel is held, opens a stream on it locally or relays through the owning replica, and hands back a conn past both handshakes; gRPC, the file stager's HTTP client and the log-streaming WebSocket are all pointed at it. A worker's address stops being somewhere to connect to and becomes the name of which backend process a stream is for. It still appears in URLs, logs and errors, because that is what identifies the process; what it no longer decides is where the bytes go. Nothing falls back to dialling it. BackendClientFactory now has exactly one method, NewClientForNode, and returns an error where there is no way to reach the worker. The direct-dial constructor was removed rather than kept beside it, because leaving one on the interface keeps the bypass one word away from every call site that holds an address, which is all of them. The second construction path is closed too. DistributedModelStore built remote models with a nil client, and pkg/model.Model.GRPC then dialled the raw address lazily on first use - reached in production by ShutdownModel's Free and by the backend monitor's Status. Those models now carry the tunnel-backed client, and a model that cannot be given one is logged and not listed. Four conditions stay unmixable, and one path produces absence: the dialer answers ErrNoConnection only where Owner's liveness join did. A peer that will not answer, a stale ownership row, a worker's own refusal and a missing relay path are each reported as themselves. This matters because nodes ACTS on absence, and the collapse would have it reclaim the models of a worker that is connected and busy. That is not hypothetical. Writing the mutation for it exposed the bug in this change's own first draft: probeHealth returned bare false when it could not build a client, and tryWarmPath deletes the replica row on a false probe. A frontend whose dialer broke would have emptied node_models for the whole deployment while every model kept running. probeHealth now returns alive and probed separately, the reconciler gets a ProbeUnknown outcome that neither advances nor clears a failure streak, and the health monitor skips rather than counting a miss. Task 5 left the relay's open timeout at a fixed 15s and said so: no operator has the information to set it, because the number that matters is the original client's remaining budget, which is invisible on the relay side. The dialer has that budget, so it now states it in the relay request frame and the owner takes the smaller of the two. It can only shorten - a patient client must not be able to park a relay goroutine and a stream slot on a worker that stopped accepting. Zero is written as no budget at all, since on the far side the number zero is a caller with nothing left and would refuse healthy traffic. Seven mutations, each reddening a named spec: peer-unreachable as absence; the local-failure guard dropped; max instead of min on the budget; the nil-client model restored; ProbeUnknown falling through to the reaper; OwnerRow instead of Owner; probed collapsed into alive. The first budget spec passed for the wrong reason - a handshake deadline, not the relay - and was replaced by three that each assert one link, including one where the spec plays the owning replica and reads the budget out of the frame instead of inferring it from a clock. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 73 ++- core/application/startup.go | 6 + core/http/app.go | 13 +- core/http/endpoints/localai/nodes.go | 53 +- core/http/routes/nodes.go | 14 +- core/services/cluster/dialer.go | 250 ++++++++ core/services/cluster/dialer_test.go | 563 ++++++++++++++++++ core/services/cluster/relay.go | 110 +++- core/services/cluster/relay_internal_test.go | 79 ++- core/services/cluster/relay_test.go | 8 +- .../nodes/backend_client_factory_test.go | 101 ++++ core/services/nodes/distributed_store.go | 48 +- core/services/nodes/distributed_store_test.go | 69 ++- core/services/nodes/file_stager_http.go | 123 +++- .../nodes/file_stager_verify_deadline_test.go | 2 +- .../nodes/file_transfer_server_test.go | 56 +- core/services/nodes/health.go | 19 +- core/services/nodes/health_mock_test.go | 23 + core/services/nodes/health_test.go | 40 ++ core/services/nodes/interfaces.go | 95 ++- .../nodes/local_stub_invalidator_test.go | 6 +- core/services/nodes/reconciler.go | 66 +- .../nodes/reconciler_busy_probe_test.go | 42 ++ core/services/nodes/reconciler_test.go | 2 +- .../nodes/revision_eligibility_test.go | 2 +- core/services/nodes/router.go | 80 ++- .../services/nodes/router_load_budget_test.go | 4 + .../nodes/router_load_timeout_test.go | 4 + core/services/nodes/router_reap_load_test.go | 4 + core/services/nodes/router_test.go | 4 + .../nodes/router_unreachable_worker_test.go | 70 +++ docs/content/features/distributed-mode.md | 31 + pkg/grpc/backend.go | 26 +- pkg/grpc/client.go | 14 + tests/e2e/distributed/backend_logs_test.go | 19 +- .../distributed/distributed_full_flow_test.go | 12 +- .../e2e/distributed/distributed_store_test.go | 22 +- tests/e2e/distributed/file_staging_test.go | 2 +- .../distributed/prefix_cache_routing_test.go | 4 + 39 files changed, 1990 insertions(+), 169 deletions(-) create mode 100644 core/services/cluster/dialer.go create mode 100644 core/services/cluster/dialer_test.go create mode 100644 core/services/nodes/backend_client_factory_test.go create mode 100644 core/services/nodes/router_unreachable_worker_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 074d261d912b..512b81780296 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -67,6 +67,14 @@ type DistributedServices struct { // loop, which re-claims what it holds after this replica has been reaped, // and to the route that accepts a worker's dial. Tunnels *cluster.TunnelRegistry + // WorkerDialer is how anything in this process reaches a worker: locally + // when this replica holds the tunnel, and through the owning replica when + // it does not. The HTTP layer takes its WebSocket log proxy from here. + WorkerDialer *cluster.WorkerDialer + // BackendClients builds the gRPC clients for worker backend processes, over + // WorkerDialer. Exposed so the model store built in startup.go reaches + // remote models the same way every other caller does. + BackendClients nodes.BackendClientFactory shutdownOnce sync.Once } @@ -254,6 +262,22 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // peer route checks (see RegisterClusterRoutes); two different tokens here // would make every peer dial 401 with nothing naming the mismatch. peers := cluster.NewPeerPool(cfg.Distributed.InstanceID, cfg.Distributed.RegistrationToken, clusterRegistry) + // The one door to every worker. Nothing in the frontend may dial a worker's + // advertised address any more: a worker holds ONE tunnel, it lands on ONE + // replica, and this resolves which replica that is and relays through it + // when it is not this one. The three transports the frontend speaks to a + // worker (gRPC to backend processes, HTTP for file staging and logs, a + // WebSocket for live log streaming) are all pointed at it below. + workerDialer := cluster.NewWorkerDialer(tunnels, peers) + backendClients, err := nodes.NewTunnelClientFactory(cfg.Distributed.RegistrationToken, workerDialer.GRPCDialerFor) + if err != nil { + return nil, fmt.Errorf("wiring the worker backend client factory: %w", err) + } + // Bound to the http tag: the worker ignores the target for it and routes to + // its own file-transfer and log server, wherever that bound. + workerHTTPDialer := nodes.WorkerNetDialerFor(func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) { + return workerDialer.DialerFor(nodeID, cluster.StreamTagHTTP) + }) // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an @@ -295,6 +319,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade cfg.Distributed.StaleNodeThresholdOrDefault(), routerAuthToken, !cfg.Distributed.DisablePerModelHealthCheck, + backendClients, ) // Initialize job store @@ -354,7 +379,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) } return node.HTTPAddress, nil - }, cfg.Distributed.RegistrationToken) + }, cfg.Distributed.RegistrationToken, workerHTTPDialer) xlog.Info("File stager initialized (HTTP direct transfer)") } // Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go @@ -456,6 +481,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade FileStager: fileStager, GalleriesJSON: routerGalleriesJSON, AuthToken: routerAuthToken, + ClientFactory: backendClients, DB: authDB, ConflictResolver: conflictResolver, PrefixProvider: prefixProvider, @@ -514,6 +540,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Unloader: remoteUnloader, Adapter: remoteUnloader, RegistrationToken: cfg.Distributed.RegistrationToken, + ClientFactory: backendClients, DB: authDB, Interval: 30 * time.Second, ScaleDownDelay: 5 * time.Minute, @@ -527,27 +554,29 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade success = true return &DistributedServices{ - Nats: natsClient, - Store: store, - Registry: registry, - Router: router, - Health: healthMon, - Reconciler: reconciler, - JobStore: jobStore, - Dispatcher: dispatcher, - AgentStore: agentStore, - AgentBridge: agentBridge, - DistStores: distStores, - FileMgr: fileMgr, - FileStager: fileStager, - ModelAdapter: modelAdapter, - Unloader: remoteUnloader, - ModelCleanup: modelCleanup, - Cluster: clusterRegistry, - Membership: membership, - PeerSessions: peerSessions, - Peers: peers, - Tunnels: tunnels, + Nats: natsClient, + Store: store, + Registry: registry, + Router: router, + Health: healthMon, + Reconciler: reconciler, + JobStore: jobStore, + Dispatcher: dispatcher, + AgentStore: agentStore, + AgentBridge: agentBridge, + DistStores: distStores, + FileMgr: fileMgr, + FileStager: fileStager, + ModelAdapter: modelAdapter, + Unloader: remoteUnloader, + ModelCleanup: modelCleanup, + Cluster: clusterRegistry, + Membership: membership, + PeerSessions: peerSessions, + Peers: peers, + Tunnels: tunnels, + WorkerDialer: workerDialer, + BackendClients: backendClients, }, nil } diff --git a/core/application/startup.go b/core/application/startup.go index abc2f4a17571..c790a1403617 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -283,9 +283,15 @@ func New(opts ...config.AppOption) (*Application, error) { // Wire ModelRouter so grpcModel() delegates to SmartRouter in distributed mode application.modelLoader.SetModelRouter(distSvc.ModelAdapter.AsModelRouter()) // Wire DistributedModelStore so shutdown/list/watchdog can find remote models + // The client factory is not optional here. Without it the store builds + // remote models with no client, and pkg/model.Model.GRPC then dials the + // worker's raw address with gRPC's own dialer, which is the direct dial + // the tunnel replaces; ShutdownModel's Free and the backend monitor's + // Status both reach it. distStore := nodes.NewDistributedModelStore( model.NewInMemoryModelStore(), distSvc.Registry, + distSvc.BackendClients, ) application.modelLoader.SetModelStore(distStore) // Drop the local stub when a model's last replica leaves the registry. diff --git a/core/http/app.go b/core/http/app.go index 5d3299027580..ec2b864ee8b7 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -1,12 +1,14 @@ package http import ( + "context" "embed" "errors" "fmt" "io/fs" "math" "mime" + "net" "net/http" "os" "path/filepath" @@ -567,15 +569,24 @@ func API(application *application.Application) (*echo.Echo, error) { distCfg := application.ApplicationConfig().Distributed var registry *nodes.NodeRegistry var remoteUnloader nodes.NodeCommandSender + // How the admin log-proxy routes reach a worker's own HTTP server. Left nil + // outside distributed mode, where there are no workers and no tunnels; the + // routes then refuse rather than dialling an address directly. + var workerHTTPDialFor nodes.WorkerNetDialerFor if d := application.Distributed(); d != nil { registry = d.Registry if d.Router != nil { remoteUnloader = d.Router.Unloader() } + if d.WorkerDialer != nil { + workerHTTPDialFor = func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) { + return d.WorkerDialer.DialerFor(nodeID, clustersvc.StreamTagHTTP) + } + } } 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) + routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDialFor) // Replica-to-replica peer link. Registered only in distributed mode: in // single-node mode there are no peers, and the route authenticates with the diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index 0d3f3ebb3f92..f10cc9c83ca9 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -10,6 +10,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "net/url" "sync" @@ -791,7 +792,7 @@ func DeleteModelOnNodeEndpoint(unloader nodes.NodeCommandSender, registry *nodes // NodeBackendLogsListEndpoint proxies a request to a worker node's /v1/backend-logs // endpoint to list model IDs that have backend logs. -func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc { +func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc { return func(c echo.Context) error { ctx := c.Request().Context() nodeID := c.Param("id") @@ -804,7 +805,7 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) } - resp, err := proxyHTTPToWorker(node.HTTPAddress, "/v1/backend-logs", registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, node.HTTPAddress, "/v1/backend-logs", registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -819,7 +820,7 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken // NodeBackendLogsLinesEndpoint proxies a request to a worker node's // /v1/backend-logs/{modelId} endpoint to get buffered log lines. -func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc { +func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc { return func(c echo.Context) error { ctx := c.Request().Context() nodeID := c.Param("id") @@ -835,7 +836,7 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke } path := "/v1/backend-logs/" + url.PathEscape(modelID) - resp, err := proxyHTTPToWorker(node.HTTPAddress, path, registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, node.HTTPAddress, path, registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -850,7 +851,7 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke // NodeBackendLogsWSEndpoint proxies a WebSocket connection to a worker node's // /v1/backend-logs/{modelId}/ws endpoint for real-time log streaming. -func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc { +func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc { browserUpgrader := websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { origin := r.Header.Get("Origin") @@ -882,15 +883,28 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s return err } - // Dial the worker WebSocket + // Dial the worker WebSocket over that worker's tunnel. The URL still + // names the worker's registered address, for the Host header; the + // NetDialContext below is what decides where the connection goes. A + // missing dialer is a failure, not a direct dial: see + // nodes.ErrNoWorkerDialer. workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", node.HTTPAddress, url.PathEscape(modelID)) workerHeaders := http.Header{} if registrationToken != "" { workerHeaders.Set("Authorization", "Bearer "+registrationToken) } - workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second} - workerWS, _, err := workerDialer.Dial(workerURL, workerHeaders) + var workerDial func(ctx context.Context, network, addr string) (net.Conn, error) + if dialFor != nil { + workerDial = dialFor(nodeID) + } + if workerDial == nil { + return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, + fmt.Sprintf("cannot reach node %s: %v", nodeID, nodes.ErrNoWorkerDialer))) + } + + workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second, NetDialContext: workerDial} + workerWS, _, err := workerDialer.DialContext(ctx, workerURL, workerHeaders) if err != nil { browserWS.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "failed to connect to worker")) @@ -1347,10 +1361,25 @@ func DeleteSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { } } -// proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer token auth. -func proxyHTTPToWorker(httpAddress, path, token string) (*http.Response, error) { +// proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer +// token auth, over that worker's tunnel. +// +// The URL still names the worker's registered HTTP address, because that is +// what the Host header and every error message should say; what it no longer +// decides is where the bytes go. dialFor supplies the transport, and a nil one +// is an error rather than a fall back to connecting to httpAddress: a worker +// behind NAT has no address to connect to, and a direct dial is the bypass this +// whole change removes. +func proxyHTTPToWorker(ctx context.Context, dialFor nodes.WorkerNetDialerFor, nodeID, httpAddress, path, token string) (*http.Response, error) { + if dialFor == nil { + return nil, fmt.Errorf("reaching node %s: %w", nodeID, nodes.ErrNoWorkerDialer) + } + dial := dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("reaching node %s: %w", nodeID, nodes.ErrNoWorkerDialer) + } reqURL := fmt.Sprintf("http://%s%s", httpAddress, path) - req, err := http.NewRequest("GET", reqURL, nil) + req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil) if err != nil { return nil, err } @@ -1358,6 +1387,6 @@ func proxyHTTPToWorker(httpAddress, path, token string) (*http.Response, error) req.Header.Set("Authorization", "Bearer "+token) } - client := httpclient.NewWithTimeout(15 * time.Second) + client := httpclient.NewWithTimeout(15*time.Second, httpclient.WithTransport(&http.Transport{DialContext: dial})) return client.Do(req) } diff --git a/core/http/routes/nodes.go b/core/http/routes/nodes.go index 053d6c19cf30..a511c6b7c1fc 100644 --- a/core/http/routes/nodes.go +++ b/core/http/routes/nodes.go @@ -61,7 +61,13 @@ 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) { +// +// workerDialFor is how the log-proxy routes reach a worker's own HTTP server: +// over the tunnel that worker holds, never by connecting to the address it +// registered. It is nil outside distributed mode, and those two routes then +// answer 502 rather than dialling, because a worker with no tunnel has nothing +// for them to proxy to. +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, workerDialFor nodes.WorkerNetDialerFor) { if registry == nil { return } @@ -101,8 +107,8 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade admin.POST("/:id/models/delete", localai.DeleteModelOnNodeEndpoint(unloader, registry)) // Backend log streaming (proxied from worker HTTP server) - admin.GET("/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, registrationToken)) - admin.GET("/:id/backend-logs/:modelId", localai.NodeBackendLogsLinesEndpoint(registry, registrationToken)) + admin.GET("/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, registrationToken, workerDialFor)) + admin.GET("/:id/backend-logs/:modelId", localai.NodeBackendLogsLinesEndpoint(registry, registrationToken, workerDialFor)) // Label management admin.GET("/:id/labels", localai.GetNodeLabelsEndpoint(registry)) @@ -123,7 +129,7 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade admin.DELETE("/:id/vram-budget", localai.ResetVRAMBudgetEndpoint(registry)) // WebSocket proxy for real-time log streaming from workers - e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, registrationToken), readyMw, adminMw) + e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, registrationToken, workerDialFor), readyMw, adminMw) } // nodeTokenAuth validates the registration token for node self-service endpoints. diff --git a/core/services/cluster/dialer.go b/core/services/cluster/dialer.go new file mode 100644 index 000000000000..b4a04551a973 --- /dev/null +++ b/core/services/cluster/dialer.go @@ -0,0 +1,250 @@ +// SPDX-License-Identifier: MIT + +package cluster + +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 a spec can decide what a peer does without standing up +// a second frontend. Note that a TYPED nil (a (*PeerPool)(nil) stored here) +// would not compare equal to nil and would be called; the dialer's nil check is +// for the untyped nil a deployment with no peer mesh passes. +type PeerOpener interface { + Open(ctx context.Context, peerID string) (net.Conn, error) +} + +// ErrNoRelayPath reports that a worker's tunnel is held by ANOTHER replica and +// this one has no way to reach that replica. +// +// It is its own condition, kept out of the four this phase refuses to collapse. +// Not ErrNotOwner, because that tells a caller to resolve the owner again and +// the answer would not change. Not ErrPeerUnreachable, because no peer was +// dialled and none refused; saying otherwise would blame a replica that is +// probably fine. And above all not ErrNoConnection: the worker IS connected, +// and core/services/nodes reclaims the models of a worker it believes absent. +var ErrNoRelayPath = errors.New("cluster: this replica cannot relay to the owner of that worker") + +// dialHandshakeTimeout bounds the request/reply exchange that opens 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 session's +// keepalive killed it, which is 30 seconds on the yamux default and longer if a +// deployment ever raises it. Where the caller DOES carry a deadline, that +// deadline is used instead whenever it is the shorter of the two, for the same +// reason the relay takes the smaller of its ceiling and the stated budget. +const dialHandshakeTimeout = 15 * time.Second + +// WorkerDialer opens connections to a worker through its tunnel, wherever in +// the deployment that tunnel happens to be held. +// +// It is the single door: a worker holds ONE tunnel, it lands on ONE frontend +// replica, and nothing else in the frontend may dial a worker's advertised +// address. Every protocol the frontend speaks to a worker (gRPC to a backend +// process, HTTP for file staging and logs, a WebSocket for log streaming) goes +// through the functions below, because a worker behind NAT has no address to +// dial and a worker that has one must not be reached that way either: a direct +// dial works in a single-replica test and fails in production. +type WorkerDialer struct { + // Both are read off the tunnel registry rather than passed separately, so + // the identity this dialer compares owners against is by construction the + // identity the registry CLAIMS as. Two ids here would make this replica + // relay to itself for every worker it holds. + tunnels *TunnelRegistry + peers PeerOpener +} + +// NewWorkerDialer returns the dialer for the tunnels this replica holds and the +// peer links it can relay over. A nil peers means this replica cannot relay, +// which is reported as ErrNoRelayPath rather than as a worker being absent. +func NewWorkerDialer(tunnels *TunnelRegistry, peers PeerOpener) *WorkerDialer { + return &WorkerDialer{tunnels: tunnels, peers: peers} +} + +// Dial opens one stream to a local service on a worker: tag says which service +// (see StreamTagGRPC and StreamTagHTTP) and target which instance of it. +// +// The returned conn is past both handshakes and carries the tunnelled protocol +// and nothing else, with no deadline armed on it: what follows may be an +// inference that is quiet for minutes, and a deadline left over from the +// handshake would abort it. +// +// The errors are kept apart on purpose and a caller may act on them +// differently. ErrNoConnection means no live replica holds this worker's +// tunnel, which is the one answer that means the worker is absent. ErrNotOwner +// means the routing was stale and re-resolving may find it. ErrPeerUnreachable +// means a replica would not answer, ErrNoRelayPath that none could be dialled, +// and the tunnelproto sentinels that the worker itself refused. Nothing here +// ever converts one of the others into absence. +func (d *WorkerDialer) Dial(ctx context.Context, nodeID, tag, target string) (net.Conn, error) { + stream, err := d.tunnels.Open(ctx, nodeID) + if err == nil { + return d.handshake(ctx, stream, nodeID, tag, target) + } + if !errors.Is(err, ErrNotOwner) { + // The tunnel is held HERE and its session would not carry a stream. + // Reported as itself: answering ErrNotOwner would send the caller to + // resolve an owner that is this same replica, and answering absence + // would tell a scheduler to reclaim a worker that is attached. + return nil, err + } + return d.relay(ctx, nodeID, tag, target) +} + +// DialerFor returns a net.Dialer-shaped function bound to one worker and one +// local service on it. +// +// This is the shape http.Transport.DialContext and websocket.Dialer's +// NetDialContext want. The NETWORK is ignored and the ADDRESS becomes the +// stream's target, which is what makes an http.Client or a WebSocket dialler +// built on it reach the worker without either of them knowing a tunnel exists: +// the URL still names the worker's registered address, and that address travels +// as the target rather than to a socket. What the worker does with it is the +// worker's decision (the grpc tag takes the port and dials its own loopback, +// the http tag ignores it entirely), which is the property that stops a +// frontend from steering a worker's dial. +func (d *WorkerDialer) DialerFor(nodeID, tag string) 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) + } +} + +// GRPCDialerFor returns a grpc.WithContextDialer-shaped function bound to one +// worker's backend processes. gRPC's dialer takes no network argument, which is +// why this is not DialerFor's shape. +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 worker's tunnel. +func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string) (net.Conn, error) { + // Owner, never OwnerRow: the row outlives its owner by up to a liveness + // window plus a heartbeat, and dialling what the unjoined read returns + // means dialling a process that is gone and reporting the worker as + // unreachable rather than as absent. The join is what makes a dead owner + // come back as ErrNoConnection here. + owner, _, err := d.tunnels.reg.Owner(ctx, nodeID) + if err != nil { + // ErrNoConnection and database failures both pass through as + // themselves. This is the ONLY path by which this function can produce + // an absence error, and it produces it only when Owner did. + return nil, err + } + if owner == d.tunnels.selfID { + // The table names this replica and the registry above said 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, which + // would resolve the same owner and relay again. Reported as the routing + // fact so the caller re-resolves, which terminates: the row is either + // re-claimed by whoever holds the worker now, or swept. + return nil, fmt.Errorf("opening a stream to node %q: the connection row names this replica, which no longer holds the tunnel: %w", nodeID, ErrNotOwner) + } + if d.peers == nil { + return nil, fmt.Errorf("opening a stream to node %q held by replica %q: %w", nodeID, owner, ErrNoRelayPath) + } + + stream, err := d.peers.Open(ctx, owner) + if err != nil { + // Whatever the pool said, unchanged in its unwrap chain: + // ErrPeerUnreachable, ErrInstanceNotFound for an owner swept since the + // lookup above, or ErrPoolClosed while this process shuts down. None of + // them is a statement about the WORKER, and none is converted into one. + return nil, fmt.Errorf("opening a stream to node %q through replica %q: %w", nodeID, owner, err) + } + + // The caller's remaining time, stated so the owning replica can bound its + // own open by it. Only this side knows it; see relayOpenTimeout for what + // the owner falls back to without it. + if err := WriteRelayRequest(stream, nodeID, remainingBudget(ctx)); err != nil { + _ = stream.Close() + return nil, fmt.Errorf("naming node %q on a stream to replica %q: %w", nodeID, owner, err) + } + if err := ReadRelayReply(stream); err != nil { + // ReadRelayReply already separates a refusal (ErrNotOwner, + // ErrRelayUnavailable, ErrRelayRequestInvalid) from a failure to read + // one, and neither kind is ever an absence error. + _ = stream.Close() + return nil, fmt.Errorf("relaying to node %q through replica %q: %w", nodeID, owner, err) + } + return d.handshake(ctx, stream, nodeID, tag, target) +} + +// handshake names the worker-side service on a stream and waits for the +// worker's answer, leaving the stream ready for the tunnelled protocol. +// +// It owns closing the stream on every failure. A stream left open after a +// failed handshake holds a yamux slot on the tunnel for the life of the +// session, and a frontend that retries would exhaust the worker's stream +// budget rather than the worker's patience. +func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, tag, target string) (net.Conn, error) { + if deadline, ok := handshakeDeadline(ctx); ok { + if err := stream.SetDeadline(deadline); err != nil { + _ = stream.Close() + return nil, fmt.Errorf("arming the handshake deadline for node %q: %w", nodeID, err) + } + } + + if err := WriteStreamRequest(stream, tag, target); err != nil { + _ = stream.Close() + return nil, fmt.Errorf("asking node %q for %q on %q: %w", nodeID, tag, target, err) + } + if err := ReadStreamReply(stream); err != nil { + _ = stream.Close() + return nil, fmt.Errorf("opening %q on node %q: %w", tag, nodeID, err) + } + + // Cleared unconditionally rather than only when one was armed, so that this + // stays true of the stream whatever the caller's context carried. What + // follows is the caller's protocol, and its length is the caller's + // business: a request may sit quiet for minutes between tokens, and the + // handshake's deadline would end it. The session's keepalive is what still + // bounds a peer that has stopped answering. + if err := stream.SetDeadline(time.Time{}); err != nil { + _ = stream.Close() + return nil, fmt.Errorf("clearing the handshake deadline for node %q: %w", nodeID, 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 by: the caller's own +// deadline when it has one and it is the sooner, and the backstop otherwise. +func handshakeDeadline(ctx context.Context) (time.Time, bool) { + backstop := time.Now().Add(dialHandshakeTimeout) + deadline, ok := ctx.Deadline() + if !ok || deadline.After(backstop) { + return backstop, true + } + return deadline, true +} + +// remainingBudget is how long the caller is still willing to wait, or zero when +// it did not say. +// +// Zero rather than a negative number for an expired context: the frame writer +// treats zero as "not stated", and a caller that has already run out is about +// to fail on its own context anyway. Stating a negative budget would instead +// make the owning replica refuse, which is the same outcome by a longer route. +func remainingBudget(ctx context.Context) time.Duration { + deadline, ok := ctx.Deadline() + if !ok { + return 0 + } + remaining := time.Until(deadline) + if remaining <= 0 { + return 0 + } + return remaining +} diff --git a/core/services/cluster/dialer_test.go b/core/services/cluster/dialer_test.go new file mode 100644 index 000000000000..ffcfe3775e13 --- /dev/null +++ b/core/services/cluster/dialer_test.go @@ -0,0 +1,563 @@ +// SPDX-License-Identifier: MIT + +package cluster_test + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "net" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// stubPeers is a PeerOpener that hands back streams on one session, or one +// error. It stands in for the replica-to-replica link so a spec can decide what +// a peer does without standing up a second frontend. +type stubPeers struct { + sess *yamux.Session + err error +} + +func (s *stubPeers) Open(ctx context.Context, _ string) (net.Conn, error) { + if s.err != nil { + return nil, s.err + } + return s.sess.OpenStream(ctx) +} + +// dialResult carries what a Dial produced, so a spec can wait on a channel +// rather than on a clock. +type dialResult struct { + conn net.Conn + err error +} + +// dialAsync runs one Dial on its own goroutine. Dial talks to a worker that a +// spec drives by hand, so the spec has to be free to answer while the dial is +// still in flight. +func dialAsync(d *cluster.WorkerDialer, ctx context.Context, nodeID, tag, target string) chan dialResult { + done := make(chan dialResult, 1) + go func() { + defer GinkgoRecover() + conn, err := d.Dial(ctx, nodeID, tag, target) + done <- dialResult{conn: conn, err: err} + }() + return done +} + +// relayRequest is what an owning replica saw in the frame that opened a +// relayed stream. +type relayRequest struct { + nodeID string + budget time.Duration + err error +} + +// servedRequest is what a worker saw on a stream opened through the tunnel. +type servedRequest struct { + tag string + target string + stream net.Conn + err error +} + +// serveOneStream accepts one stream on the worker's half, reads the tunnel +// request frame and accepts it, then echoes the four bytes it is sent. It is +// how a spec proves the dial produced a stream that carries the tunnelled +// protocol, and that the frame the worker sees is the one the dialer wrote. +func serveOneStream(worker *yamux.Session) chan servedRequest { + seen := make(chan servedRequest, 1) + go func() { + defer GinkgoRecover() + stream, err := worker.AcceptStream() + if err != nil { + seen <- servedRequest{err: err} + return + } + tag, target, err := cluster.ReadStreamRequest(stream) + if err != nil { + seen <- servedRequest{err: err} + return + } + if err := cluster.WriteStreamAccepted(stream); err != nil { + seen <- servedRequest{err: err} + return + } + seen <- servedRequest{tag: tag, target: target, stream: stream} + buf := make([]byte, 4) + if _, err := io.ReadFull(stream, buf); err != nil { + return + } + _, _ = stream.Write(buf) + }() + return seen +} + +// refuseOneStream accepts one stream and refuses it with reason. +func refuseOneStream(worker *yamux.Session, reason error) { + go func() { + defer GinkgoRecover() + stream, err := worker.AcceptStream() + if err != nil { + return + } + defer func() { _ = stream.Close() }() + if _, _, err := cluster.ReadStreamRequest(stream); err != nil { + return + } + _ = cluster.WriteStreamRefusal(stream, reason) + }() +} + +// expectNotAbsence asserts an error is none of the sentinels a caller is +// entitled to act on as "this worker has gone away". +// +// It is the assertion this whole phase turns on. core/services/nodes reclaims a +// worker's models when it concludes the worker is absent, so an unreachable +// peer or a refusing worker arriving as absence would evict healthy work. +func expectNotAbsence(err error) { + GinkgoHelper() + Expect(err).To(HaveOccurred()) + Expect(err).ToNot(MatchError(cluster.ErrNoConnection)) + Expect(err).ToNot(MatchError(cluster.ErrInstanceNotFound)) +} + +var _ = Describe("The worker dialer", func() { + var ( + db *gorm.DB + reg *cluster.Registry + mine *cluster.TunnelRegistry + ctx context.Context + ) + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + mine = cluster.NewTunnelRegistry(reg, "me") + }) + + // ownerRelay stands up a SECOND replica that holds w1's tunnel and relays + // for it, and returns the peer opener this replica reaches it through plus + // the worker's own half of the tunnel. + ownerRelay := func(nodeID string) (*stubPeers, *yamux.Session) { + GinkgoHelper() + Expect(reg.Register(ctx, "owner", "10.0.0.2:8080", "v1")).To(Succeed()) + ownerTunnels := cluster.NewTunnelRegistry(reg, "owner") + frontend, worker := workerTunnel() + _, err := ownerTunnels.Attach(ctx, nodeID, frontend) + Expect(err).ToNot(HaveOccurred()) + + store := cluster.NewSessionStore(cluster.NewRelay(ownerTunnels).Stream) + DeferCleanup(store.CloseAll) + dialling, accepted := yamuxPair() + store.Accept("me", accepted) + return &stubPeers{sess: dialling}, worker + } + + Describe("when this replica holds the tunnel", func() { + It("opens a stream straight down it, naming the service the caller asked for", func() { + frontend, worker := workerTunnel() + _, err := mine.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + seen := serveOneStream(worker) + + d := cluster.NewWorkerDialer(mine, nil) + result := dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000") + + var req servedRequest + Eventually(seen, "10s").Should(Receive(&req)) + Expect(req.err).ToNot(HaveOccurred()) + Expect(req.tag).To(Equal(cluster.StreamTagGRPC)) + // The TARGET is what tells the worker which backend process the + // stream is for. A dialer that dropped it would send every request + // to whichever port the worker guessed. + Expect(req.target).To(Equal("127.0.0.1:41000")) + + var out dialResult + Eventually(result, "10s").Should(Receive(&out)) + Expect(out.err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = out.conn.Close() }) + + // Bytes, not a handle. A dial that returns a stream the tunnelled + // protocol cannot use is worse than one that fails. + _, err = out.conn.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + echoed := make([]byte, 4) + Eventually(readInto(out.conn, echoed), "10s").Should(Receive(BeNil())) + Expect(string(echoed)).To(Equal("ping")) + }) + + It("leaves no read deadline armed on the stream it hands back", func() { + // The handshake is bounded; the request that follows it is the + // caller's business and may be a generation that is quiet for + // minutes. A deadline left armed here would abort it. + frontend, worker := workerTunnel() + _, err := mine.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + seen := serveOneStream(worker) + + deadlined, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer cancel() + d := cluster.NewWorkerDialer(mine, nil) + result := dialAsync(d, deadlined, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000") + Eventually(seen, "10s").Should(Receive()) + + var out dialResult + Eventually(result, "10s").Should(Receive(&out)) + Expect(out.err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = out.conn.Close() }) + + // The dial context's deadline has now passed. A stream still + // carrying it would fail this write and this read. + Eventually(func() error { + _, err := out.conn.Write([]byte("ping")) + return err + }, "10s").Should(Succeed()) + echoed := make([]byte, 4) + Eventually(readInto(out.conn, echoed), "10s").Should(Receive(BeNil())) + Expect(string(echoed)).To(Equal("ping")) + }) + + It("reports a worker's refusal as the worker's refusal, never as absence", func() { + frontend, worker := workerTunnel() + _, err := mine.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + refuseOneStream(worker, cluster.ErrStreamTagUnknown) + + d := cluster.NewWorkerDialer(mine, nil) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", "nonsense", ""), "10s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrStreamTagUnknown)) + // A refusal is PROOF the worker is connected and answered. + expectNotAbsence(out.err) + }) + + It("reports a broken tunnel held here as itself, not as a routing fact", func() { + // ErrNotOwner tells a caller to look for the worker elsewhere. For + // a tunnel held right here that sends it back to this replica, and + // the loop is only broken by the request failing anyway. + frontend, worker := workerTunnel() + _, err := mine.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + Expect(worker.Close()).To(Succeed()) + + d := cluster.NewWorkerDialer(mine, nil) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + Expect(out.err).To(HaveOccurred()) + Expect(out.err).ToNot(MatchError(cluster.ErrNotOwner)) + expectNotAbsence(out.err) + }) + }) + + Describe("when another replica holds the tunnel", func() { + It("relays through the owner and carries bytes to the worker", func() { + peers, worker := ownerRelay("w1") + seen := serveOneStream(worker) + + d := cluster.NewWorkerDialer(mine, peers) + result := dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000") + + var req servedRequest + Eventually(seen, "10s").Should(Receive(&req)) + Expect(req.err).ToNot(HaveOccurred()) + // The relay consumed its own frame and forwarded nothing of it, so + // the worker sees only the tunnel's request. + Expect(req.tag).To(Equal(cluster.StreamTagGRPC)) + Expect(req.target).To(Equal("127.0.0.1:41000")) + + var out dialResult + Eventually(result, "10s").Should(Receive(&out)) + Expect(out.err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = out.conn.Close() }) + + _, err := out.conn.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + echoed := make([]byte, 4) + Eventually(readInto(out.conn, echoed), "10s").Should(Receive(BeNil())) + Expect(string(echoed)).To(Equal("ping")) + }) + + It("states the caller's remaining budget in the relay request", func() { + // The owning replica's own open bound is a backstop nobody can set + // correctly: the number that matters is how long the ORIGINAL + // client will wait, and this replica is the only one that holds it. + // This spec plays the owner by hand so it can read the frame rather + // than infer it from a timing. + Expect(reg.Register(ctx, "owner", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "owner") + Expect(err).ToNot(HaveOccurred()) + dialling, ownerSide := yamuxPair() + DeferCleanup(func() { _ = ownerSide.Close() }) + + requests := make(chan relayRequest, 1) + go func() { + defer GinkgoRecover() + stream, err := ownerSide.AcceptStream() + if err != nil { + return + } + defer func() { _ = stream.Close() }() + nodeID, budget, err := cluster.ReadRelayRequest(stream) + requests <- relayRequest{nodeID: nodeID, budget: budget, err: err} + }() + + budgeted, cancel := context.WithTimeout(ctx, 4*time.Second) + defer cancel() + d := cluster.NewWorkerDialer(mine, &stubPeers{sess: dialling}) + dialAsync(d, budgeted, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000") + + var req relayRequest + Eventually(requests, "10s").Should(Receive(&req)) + Expect(req.err).ToNot(HaveOccurred()) + Expect(req.nodeID).To(Equal("w1")) + // Whatever is left of the four seconds, and nothing invented: a + // dialer that stated its own constant would satisfy neither bound. + Expect(req.budget).To(BeNumerically(">", 2*time.Second)) + Expect(req.budget).To(BeNumerically("<=", 4*time.Second)) + }) + + It("states no budget at all for a caller that set no deadline", func() { + // Zero on the wire would be read by the owner as a caller with + // nothing left, and it would refuse traffic that is perfectly + // healthy. "Not stated" has to stay distinguishable from "expired". + Expect(reg.Register(ctx, "owner", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "owner") + Expect(err).ToNot(HaveOccurred()) + dialling, ownerSide := yamuxPair() + DeferCleanup(func() { _ = ownerSide.Close() }) + + requests := make(chan relayRequest, 1) + go func() { + defer GinkgoRecover() + stream, err := ownerSide.AcceptStream() + if err != nil { + return + } + defer func() { _ = stream.Close() }() + nodeID, budget, err := cluster.ReadRelayRequest(stream) + requests <- relayRequest{nodeID: nodeID, budget: budget, err: err} + }() + + d := cluster.NewWorkerDialer(mine, &stubPeers{sess: dialling}) + dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000") + + var req relayRequest + Eventually(requests, "10s").Should(Receive(&req)) + Expect(req.err).ToNot(HaveOccurred()) + Expect(req.budget).To(BeZero()) + }) + + It("reports an unreachable peer as unreachable, NEVER as absence", func() { + // The catastrophe this phase exists to prevent. A scheduler ACTS on + // absence: told a connected worker is gone, it reclaims every model + // the worker is running. + Expect(reg.Register(ctx, "owner", "127.0.0.1:1", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "owner") + Expect(err).ToNot(HaveOccurred()) + + pool := cluster.NewPeerPool("me", "tok", reg) + DeferCleanup(pool.Close) + d := cluster.NewWorkerDialer(mine, pool) + + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "20s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrPeerUnreachable)) + expectNotAbsence(out.err) + }) + + It("passes a stale ownership refusal back as the routing fact", func() { + // The owner's table row survives a tunnel that has gone. The relay + // answers ErrNotOwner, and only that answer tells this replica to + // resolve the owner again rather than give up on the worker. + Expect(reg.Register(ctx, "owner", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "owner") + Expect(err).ToNot(HaveOccurred()) + ownerTunnels := cluster.NewTunnelRegistry(reg, "owner") + store := cluster.NewSessionStore(cluster.NewRelay(ownerTunnels).Stream) + DeferCleanup(store.CloseAll) + dialling, accepted := yamuxPair() + store.Accept("me", accepted) + + d := cluster.NewWorkerDialer(mine, &stubPeers{sess: dialling}) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrNotOwner)) + expectNotAbsence(out.err) + }) + + It("refuses rather than relaying to itself when the table names this replica", func() { + // The row says this replica owns the tunnel and the registry says + // it does not hold it. Relaying would send the request to this same + // process, which would resolve the same owner and relay again. + _, err := reg.Claim(ctx, "w1", "me") + Expect(err).ToNot(HaveOccurred()) + + d := cluster.NewWorkerDialer(mine, &stubPeers{err: errors.New("no peer should be dialled")}) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrNotOwner)) + expectNotAbsence(out.err) + }) + + It("reports having no way to relay as its own condition", func() { + Expect(reg.Register(ctx, "owner", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "owner") + Expect(err).ToNot(HaveOccurred()) + + d := cluster.NewWorkerDialer(mine, nil) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrNoRelayPath)) + Expect(out.err).ToNot(MatchError(cluster.ErrNotOwner)) + Expect(out.err).ToNot(MatchError(cluster.ErrPeerUnreachable)) + expectNotAbsence(out.err) + }) + }) + + Describe("when no live replica holds the tunnel", func() { + It("reports absence when the worker has no connection row at all", func() { + d := cluster.NewWorkerDialer(mine, &stubPeers{err: errors.New("no peer should be dialled")}) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("reports absence when the row's owner has stopped heartbeating", func() { + // End to end over the join Owner does: the row is there, the owner + // is not. A dialer built on the unjoined read would dial a corpse + // and report the worker as unreachable rather than as absent. + Expect(reg.Register(ctx, "ghost", "10.0.0.9:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "ghost") + Expect(err).ToNot(HaveOccurred()) + Expect(db.Exec( + `UPDATE instances SET last_seen = now() - make_interval(secs => ?) WHERE id = ?`, + cluster.InstanceLiveness.Seconds()*4, "ghost").Error).To(Succeed()) + + owner, _, err := reg.OwnerRow(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("ghost")) + + d := cluster.NewWorkerDialer(mine, &stubPeers{err: errors.New("no peer should be dialled")}) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + Expect(out.err).To(MatchError(cluster.ErrNoConnection)) + }) + }) + + Describe("the dialer functions it hands to the transports", func() { + It("binds one node and one tag, and passes the address through as the target", func() { + frontend, worker := workerTunnel() + _, err := mine.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + seen := serveOneStream(worker) + + d := cluster.NewWorkerDialer(mine, nil) + dial := d.DialerFor("w1", cluster.StreamTagHTTP) + done := make(chan dialResult, 1) + go func() { + defer GinkgoRecover() + conn, err := dial(ctx, "tcp", "10.0.0.3:9090") + done <- dialResult{conn: conn, err: err} + }() + + var req servedRequest + Eventually(seen, "10s").Should(Receive(&req)) + Expect(req.err).ToNot(HaveOccurred()) + Expect(req.tag).To(Equal(cluster.StreamTagHTTP)) + Expect(req.target).To(Equal("10.0.0.3:9090")) + + var out dialResult + Eventually(done, "10s").Should(Receive(&out)) + Expect(out.err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = out.conn.Close() }) + }) + + It("gives gRPC a dialer fixed on the grpc tag", func() { + frontend, worker := workerTunnel() + _, err := mine.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + seen := serveOneStream(worker) + + d := cluster.NewWorkerDialer(mine, nil) + dial := d.GRPCDialerFor("w1") + done := make(chan dialResult, 1) + go func() { + defer GinkgoRecover() + conn, err := dial(ctx, "127.0.0.1:41000") + done <- dialResult{conn: conn, err: err} + }() + + var req servedRequest + Eventually(seen, "10s").Should(Receive(&req)) + Expect(req.err).ToNot(HaveOccurred()) + Expect(req.tag).To(Equal(cluster.StreamTagGRPC)) + Expect(req.target).To(Equal("127.0.0.1:41000")) + + var out dialResult + Eventually(done, "10s").Should(Receive(&out)) + Expect(out.err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = out.conn.Close() }) + }) + }) +}) + +var _ = Describe("The relay request frame", func() { + It("carries a stated budget and reads it back", func() { + frame := &bytes.Buffer{} + Expect(cluster.WriteRelayRequest(frame, "node-7", 2500*time.Millisecond)).To(Succeed()) + nodeID, budget, err := cluster.ReadRelayRequest(frame) + Expect(err).ToNot(HaveOccurred()) + Expect(nodeID).To(Equal("node-7")) + Expect(budget).To(Equal(2500 * time.Millisecond)) + }) + + It("writes no budget at all when none is stated", func() { + // Zero must not reach the wire as the number zero: on the far side that + // is a caller with no time left, and the relay would refuse healthy + // traffic instead of falling back to its ceiling. + frame := &bytes.Buffer{} + Expect(cluster.WriteRelayRequest(frame, "node-7", 0)).To(Succeed()) + Expect(frame.Len()).To(Equal(2 + len("node-7"))) + nodeID, budget, err := cluster.ReadRelayRequest(frame) + Expect(err).ToNot(HaveOccurred()) + Expect(nodeID).To(Equal("node-7")) + Expect(budget).To(BeZero()) + }) + + It("refuses a node id that would split across the separator", func() { + Expect(cluster.WriteRelayRequest(&bytes.Buffer{}, "node 7", time.Second)).ToNot(Succeed()) + Expect(cluster.WriteRelayRequest(&bytes.Buffer{}, "", time.Second)).ToNot(Succeed()) + }) + + It("rejects a budget that is not a number of milliseconds", func() { + // The writer cannot produce this; a mismatched peer can, and reading it + // as "not stated" would silently restore the ceiling this frame exists + // to replace. + var raw bytes.Buffer + writeRawFrame(&raw, "node-7 soon") + _, _, err := cluster.ReadRelayRequest(&raw) + Expect(err).To(HaveOccurred()) + }) + + It("treats an expired stated budget as an error rather than as silence", func() { + var raw bytes.Buffer + writeRawFrame(&raw, "node-7 0") + _, _, err := cluster.ReadRelayRequest(&raw) + Expect(err).To(HaveOccurred()) + Expect(fmt.Sprint(err)).To(ContainSubstring("expired")) + }) +}) diff --git a/core/services/cluster/relay.go b/core/services/cluster/relay.go index 9fa03807d465..9392834f4e14 100644 --- a/core/services/cluster/relay.go +++ b/core/services/cluster/relay.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "net" + "strconv" "strings" "time" @@ -59,35 +60,78 @@ var ( ErrRelayRequestInvalid = errors.New("cluster: the owning replica rejected the relay request as malformed") ) -// WriteRelayRequest names the worker a peer stream is for. +// WriteRelayRequest names the worker a peer stream is for, and how much time +// the ORIGINAL client still has. +// +// The budget is what makes the relay's own open bound honest. Everything on the +// far side of this frame is work done on behalf of a caller the relay cannot +// see, so without it the relay can only fall back to a deployment-wide constant +// that no operator has the information to set (see relayOpenTimeout). The +// dialler does have the information, because it holds the caller's context, so +// it is the one that states it. +// +// A zero budget means "not stated" and is written as no budget at all, which is +// also what an older replica sends. It is NEVER written as the number zero: on +// the far side that would be indistinguishable from a caller with no time left, +// and the relay would refuse traffic that is perfectly healthy. // // An empty node id is refused here rather than on the wire, so a caller with a -// bug learns at once instead of a round trip later. -func WriteRelayRequest(w io.Writer, nodeID string) error { +// bug learns at once instead of a round trip later. So is a node id containing +// the separator, because the split below takes the FIRST one and a node id with +// a space in it would silently move part of itself into the budget. +func WriteRelayRequest(w io.Writer, nodeID string, budget time.Duration) error { if nodeID == "" { return fmt.Errorf("writing a relay request: empty node id") } - return writeFrame(w, nodeID) + if strings.Contains(nodeID, streamRequestSeparator) { + return fmt.Errorf("writing a relay request: node id %q contains a space", nodeID) + } + if budget <= 0 { + return writeFrame(w, nodeID) + } + // A plain count of milliseconds rather than a duration string: an integer + // has one spelling, so two replicas cannot disagree about it the way they + // could about a units vocabulary that grew between their versions. Rounded + // UP so a sub-millisecond budget stays positive and keeps meaning "almost + // none" rather than collapsing into "not stated". + millis := (budget + time.Millisecond - 1) / time.Millisecond + return writeFrame(w, nodeID+streamRequestSeparator+strconv.FormatInt(int64(millis), 10)) } -// ReadRelayRequest reads the opening frame of a peer stream. +// ReadRelayRequest reads the opening frame of a peer stream. The budget is zero +// when the dialling replica stated none, which is also what a replica too old +// to state one sends. // // A malformed frame is an ordinary error, NOT ErrRelayRequestInvalid: that // sentinel is what a relay SENDS to describe a refusal, and producing it here // would leave a caller unable to tell "the peer refused my request" from "I // could not read the peer's". -func ReadRelayRequest(r io.Reader) (string, error) { +func ReadRelayRequest(r io.Reader) (string, time.Duration, error) { payload, err := readFrame(r) if err != nil { - return "", fmt.Errorf("reading a relay request: %w", err) + return "", 0, fmt.Errorf("reading a relay request: %w", err) } - if payload == "" { + nodeID, budgetText, stated := strings.Cut(payload, streamRequestSeparator) + if nodeID == "" { // An empty payload is a well-formed frame naming no worker. Treating // it as a node called "" would send the caller a routing refusal for a // request no replica can ever serve, so it stays the caller's bug. - return "", fmt.Errorf("reading a relay request: empty node id") + return "", 0, fmt.Errorf("reading a relay request: empty node id") + } + if !stated { + return nodeID, 0, nil } - return payload, nil + millis, err := strconv.ParseInt(budgetText, 10, 64) + if err != nil { + return "", 0, fmt.Errorf("reading a relay request for node %q: budget %q is not a number of milliseconds: %w", nodeID, budgetText, err) + } + if millis <= 0 { + // A caller with nothing left to spend. Reported as such rather than + // folded into "not stated", so the relay refuses at once instead of + // waiting out a backstop on behalf of a client that has already gone. + return nodeID, 0, fmt.Errorf("reading a relay request for node %q: budget %d ms has already expired", nodeID, millis) + } + return nodeID, time.Duration(millis) * time.Millisecond, nil } // WriteRelayAccepted tells the peer the stream now carries the worker tunnel's @@ -172,21 +216,25 @@ const ( // peer link dies, which is minutes on the default keepalive. relayHeaderTimeout = 15 * time.Second - // relayOpenTimeout bounds opening the worker-side stream. yamux blocks an - // Open once AcceptBacklog SYNs are in flight, waiting on synCh rather than - // failing (go-yamux/v5@v5.1.0/session.go:205-212); it honours the context, - // which is the only reason there is one here. Without the bound, a worker - // that has stopped accepting would turn a refusable condition into a parked - // peer, which is the one outcome this path exists to avoid. + // relayOpenTimeout is the CEILING on opening the worker-side stream. yamux + // blocks an Open once AcceptBacklog SYNs are in flight, waiting on synCh + // rather than failing (go-yamux/v5@v5.1.0/session.go:205-212); it honours + // the context, which is the only reason there is one here. Without the + // bound, a worker that has stopped accepting would turn a refusable + // condition into a parked peer, which is the one outcome this path exists + // to avoid. // - // It is deliberately NOT configurable. No operator has the information to - // set it: the number that matters is how long the ORIGINAL client is - // willing to wait, which is not known on this side of the link and is not - // something a deployment-wide constant can stand in for. The honest fix is - // the caller's remaining budget travelling in the relay request frame, and - // that belongs to the dialler that has the budget. Until then this is a - // backstop against parking, generous on purpose, because refusing healthy - // traffic costs more than waiting. + // It is deliberately NOT configurable, and it is no longer the whole + // answer. The number that actually matters is how long the ORIGINAL client + // is willing to wait, which no deployment-wide constant can stand in for; + // the dialling replica now states it in the request frame and accept takes + // the SMALLER of the two. This remains the backstop for a caller that + // stated nothing, generous on purpose, because refusing healthy traffic + // costs more than waiting. + // + // The stated budget only ever SHORTENS the wait. A caller willing to wait + // an hour must not be able to park this replica's relay goroutine and a + // yamux stream slot for an hour on a worker that has stopped accepting. relayOpenTimeout = 15 * time.Second ) @@ -271,7 +319,7 @@ func (r *Relay) accept(peerID string, stream net.Conn) (net.Conn, bool) { return nil, false } - nodeID, err := ReadRelayRequest(stream) + nodeID, budget, err := ReadRelayRequest(stream) if err != nil { // Includes the deadline above expiring. Both are "this stream never // said which worker it wanted", which is the dialling replica's bug @@ -296,7 +344,17 @@ func (r *Relay) accept(peerID string, stream net.Conn) (net.Conn, bool) { // Not the peer's deadline, because there is none to inherit: a yamux stream // carries no context. This bounds only the open, so a request that gets // past it is never cut short by it. - ctx, cancel := context.WithTimeout(context.Background(), r.openTimeout) + // + // The SMALLER of the ceiling and what the caller said it still has. Taking + // the caller's number when it is larger would let one patient client park a + // relay goroutine on a worker that has stopped accepting for as long as it + // liked; taking the ceiling when the caller's is smaller would keep waiting + // on behalf of a client that has already given up. + open := r.openTimeout + if budget > 0 && budget < open { + open = budget + } + ctx, cancel := context.WithTimeout(context.Background(), open) defer cancel() local, err := r.tunnels.Open(ctx, nodeID) diff --git a/core/services/cluster/relay_internal_test.go b/core/services/cluster/relay_internal_test.go index e23f6c55c888..2a5e631b0ccd 100644 --- a/core/services/cluster/relay_internal_test.go +++ b/core/services/cluster/relay_internal_test.go @@ -105,7 +105,7 @@ type unwritableStream struct { func newUnwritableStream(nodeID string) *unwritableStream { GinkgoHelper() frame := &bytes.Buffer{} - Expect(WriteRelayRequest(frame, nodeID)).To(Succeed()) + Expect(WriteRelayRequest(frame, nodeID, 0)).To(Succeed()) return &unwritableStream{request: frame.Bytes(), closed: make(chan struct{})} } @@ -176,7 +176,7 @@ var _ = Describe("The relay's own budgets", func() { stream, err := peer.OpenStream(ctx) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = stream.Close() }) - Expect(WriteRelayRequest(stream, "w1")).To(Succeed()) + Expect(WriteRelayRequest(stream, "w1", 0)).To(Succeed()) replies := make(chan error, 1) go func() { @@ -243,7 +243,7 @@ var _ = Describe("The relay's own budgets", func() { stream, err := peer.OpenStream(ctx) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = stream.Close() }) - Expect(WriteRelayRequest(stream, "w1")).To(Succeed()) + Expect(WriteRelayRequest(stream, "w1", 0)).To(Succeed()) replies := make(chan error, 1) go func() { @@ -265,6 +265,79 @@ var _ = Describe("The relay's own budgets", func() { Eventually(ends, "10s").Should(Receive(HaveOccurred())) }) + It("bounds its open by the caller's stated budget when that is the shorter", func() { + // The ceiling here is 10s and the caller says it has 50ms. Without the + // stated budget this replica would hold a relay goroutine and a yamux + // stream slot for the full ceiling on behalf of a client that gave up + // almost immediately. + relay := newRelay(tun, 0, 10*time.Second) + store := NewSessionStore(relay.Stream) + DeferCleanup(store.CloseAll) + peer, accepted := backloggedPair(256) + store.Accept("peer-1", accepted) + + _, frontend := backloggedPair(1) + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + // One un-accepted open fills the one-deep backlog, so the relay's own + // open is the one that has to wait out a budget. + filler, err := frontend.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = filler.Close() }) + + stream, err := peer.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + Expect(WriteRelayRequest(stream, "w1", 50*time.Millisecond)).To(Succeed()) + + replies := make(chan error, 1) + go func() { + defer GinkgoRecover() + replies <- ReadRelayReply(stream) + }() + var reply error + // Two seconds is twenty times the stated budget and a fifth of the + // ceiling, so only a relay that honoured the budget answers inside it. + Eventually(replies, "2s").Should(Receive(&reply)) + Expect(reply).To(MatchError(ErrRelayUnavailable)) + // The tunnel IS held here. A budget running out must not turn into a + // routing fact, and it must never become absence. + Expect(reply).ToNot(MatchError(ErrNotOwner)) + Expect(reply).ToNot(MatchError(ErrNoConnection)) + }) + + It("does not let a stated budget stretch its own ceiling", func() { + // A patient client must not be able to park this replica. The ceiling + // is 50ms and the caller says it will wait ten seconds; the refusal + // still has to arrive on the ceiling. + relay := newRelay(tun, 0, 50*time.Millisecond) + store := NewSessionStore(relay.Stream) + DeferCleanup(store.CloseAll) + peer, accepted := backloggedPair(256) + store.Accept("peer-1", accepted) + + _, frontend := backloggedPair(1) + _, err := tun.Attach(ctx, "w1", frontend) + Expect(err).ToNot(HaveOccurred()) + filler, err := frontend.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = filler.Close() }) + + stream, err := peer.OpenStream(ctx) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + Expect(WriteRelayRequest(stream, "w1", 10*time.Second)).To(Succeed()) + + replies := make(chan error, 1) + go func() { + defer GinkgoRecover() + replies <- ReadRelayReply(stream) + }() + var reply error + Eventually(replies, "2s").Should(Receive(&reply)) + Expect(reply).To(MatchError(ErrRelayUnavailable)) + }) + It("closes the worker's stream when it cannot tell the peer the stream was accepted", func() { // The reply is the last thing that can fail after a worker stream has // been opened. A relay that gave up without closing it would leak one diff --git a/core/services/cluster/relay_test.go b/core/services/cluster/relay_test.go index e0c7bd0499c2..04fa46f20285 100644 --- a/core/services/cluster/relay_test.go +++ b/core/services/cluster/relay_test.go @@ -89,7 +89,7 @@ var _ = Describe("The inter-replica relay", func() { stream, err := peer.OpenStream(ctx) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = stream.Close() }) - Expect(cluster.WriteRelayRequest(stream, nodeID)).To(Succeed()) + Expect(cluster.WriteRelayRequest(stream, nodeID, 0)).To(Succeed()) return stream } @@ -279,14 +279,14 @@ var _ = Describe("The relay wire framing", func() { It("round-trips a node id", func() { frame := &bytes.Buffer{} - Expect(cluster.WriteRelayRequest(frame, "node-7")).To(Succeed()) - nodeID, err := cluster.ReadRelayRequest(frame) + Expect(cluster.WriteRelayRequest(frame, "node-7", 0)).To(Succeed()) + nodeID, _, err := cluster.ReadRelayRequest(frame) Expect(err).ToNot(HaveOccurred()) Expect(nodeID).To(Equal("node-7")) }) It("refuses to write an empty node id, rather than spending a round trip on it", func() { - Expect(cluster.WriteRelayRequest(&bytes.Buffer{}, "")).To(HaveOccurred()) + Expect(cluster.WriteRelayRequest(&bytes.Buffer{}, "", 0)).To(HaveOccurred()) }) // Acceptance is not the whole surface. A refusal read by the wrong hop's diff --git a/core/services/nodes/backend_client_factory_test.go b/core/services/nodes/backend_client_factory_test.go new file mode 100644 index 000000000000..be20f5e347b0 --- /dev/null +++ b/core/services/nodes/backend_client_factory_test.go @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpcpkg "github.com/mudler/LocalAI/pkg/grpc" +) + +var _ = Describe("The backend client factory", func() { + Describe("without a worker tunnel dialer", func() { + It("refuses to build a client for a node rather than dialling its address", func() { + // The whole point. A factory that answered here with a client + // pointed at the raw address would work on a single-host developer + // setup and fail against every worker that has no inbound port, + // which is the worst way for this to behave. + f := &tokenClientFactory{token: "tok"} + _, err := f.NewClientForNode("node-1", "10.0.0.1:41000", false) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + + It("offers no direct-dial constructor for anything to reach for", func() { + // Structural, not documented. A NewClient alongside NewClientForNode + // would be reachable from every call site that holds an address, + // which is all of them, and reintroducing the bypass would then be + // a one-word edit that compiles and passes every other spec. + var factory any = &tunnelClientFactory{} + _, hasDirectDial := factory.(interface { + NewClient(address string, parallel bool) grpcpkg.Backend + }) + Expect(hasDirectDial).To(BeFalse()) + }) + + It("refuses to be constructed at all", func() { + _, err := NewTunnelClientFactory("tok", nil) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + }) + + Describe("with a worker tunnel dialer", func() { + It("builds a client that reaches the backend through the node's dialer", func() { + // The proof is that the client's transport is the one this factory + // was given: it carries bytes from a listener that the address in + // the request never names. + listener, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = listener.Close() }) + + asked := make(chan string, 4) + f, err := NewTunnelClientFactory("", func(nodeID string) func(ctx context.Context, addr string) (net.Conn, error) { + return func(ctx context.Context, addr string) (net.Conn, error) { + asked <- nodeID + "|" + addr + var d net.Dialer + return d.DialContext(ctx, "tcp", listener.Addr().String()) + } + }) + Expect(err).ToNot(HaveOccurred()) + + client, err := f.NewClientForNode("node-1", "10.255.255.1:41000", false) + Expect(err).ToNot(HaveOccurred()) + + // The address is unroutable on purpose: only a client that used the + // dialer can reach anything at all. The health check itself fails, + // because nothing on the far side speaks gRPC; what it proves is + // which transport was asked. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + defer GinkgoRecover() + _, _ = client.HealthCheck(ctx) + }() + Eventually(asked, "10s").Should(Receive(Equal("node-1|10.255.255.1:41000"))) + }) + + It("refuses a request with no node id", func() { + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + var d net.Dialer + return func(ctx context.Context, addr string) (net.Conn, error) { + return d.DialContext(ctx, "tcp", addr) + } + }) + Expect(err).ToNot(HaveOccurred()) + _, err = f.NewClientForNode("", "10.0.0.1:41000", false) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + + It("refuses when the dialer has none for that node", func() { + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + return nil + }) + Expect(err).ToNot(HaveOccurred()) + _, err = f.NewClientForNode("node-1", "10.0.0.1:41000", false) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + }) +}) diff --git a/core/services/nodes/distributed_store.go b/core/services/nodes/distributed_store.go index ba1379367413..542aec69e693 100644 --- a/core/services/nodes/distributed_store.go +++ b/core/services/nodes/distributed_store.go @@ -2,7 +2,9 @@ package nodes import ( "context" + "fmt" + grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/xlog" ) @@ -14,10 +16,26 @@ import ( type DistributedModelStore struct { local model.ModelStore registry ModelLookup + // clients builds the gRPC client for a model that lives on a worker. + // + // It is not optional in a real deployment, and the reason is the second + // construction path this store used to be: a *model.Model built with a nil + // client makes pkg/model.Model.GRPC dial its raw address with gRPC's own + // dialer the first time anything touches it, which is exactly the direct + // dial to a worker's advertised address the tunnel replaces. That path is + // reached in production, by ShutdownModel's Free and by the backend + // monitor's Status, so it is not theoretical. + clients BackendClientFactory } -func NewDistributedModelStore(local model.ModelStore, registry ModelLookup) *DistributedModelStore { - return &DistributedModelStore{local: local, registry: registry} +// NewDistributedModelStore returns the store, which reaches a remote model's +// backend through clients. +// +// A nil clients is a programming error and is treated as one: Range refuses to +// synthesise a model it cannot give a working client to, rather than handing +// back one that silently dials the worker's address. See the field comment. +func NewDistributedModelStore(local model.ModelStore, registry ModelLookup, clients BackendClientFactory) *DistributedModelStore { + return &DistributedModelStore{local: local, registry: registry, clients: clients} } // Get checks the local cache only. In distributed mode, models must be routed @@ -80,9 +98,33 @@ func (s *DistributedModelStore) Range(fn func(string, *model.Model) bool) { continue } - m := model.NewModel(nm.ModelName, node.Address, nil) + // NewModelWithClient, never NewModel: a model built without a client + // lazily dials its address with gRPC's default dialer the first time + // anything calls GRPC() on it, which reaches a worker only while + // workers still listen on a routable address. Building the client here + // means the bypass has no path left rather than an unused one. + client, err := s.clientFor(nm.NodeID, node.Address) + if err != nil { + xlog.Error("DistributedModelStore: not listing a remote model it cannot reach", + "model", nm.ModelName, "nodeID", nm.NodeID, "error", err) + continue + } + m := model.NewModelWithClient(nm.ModelName, node.Address, client) if !fn(nm.ModelName, m) { return } } } + +// clientFor builds the backend client for a model running on a worker. +// +// It fails rather than falling back. A fallback here would be invisible: the +// listing would look complete, shutdown would appear to work, and the direct +// dial underneath it would succeed on a single-host developer setup and fail +// against every worker that has no inbound port. +func (s *DistributedModelStore) clientFor(nodeID, address string) (grpc.Backend, error) { + if s.clients == nil { + return nil, fmt.Errorf("no backend client factory is wired into the distributed model store: %w", ErrNoWorkerDialer) + } + return s.clients.NewClientForNode(nodeID, address, false) +} diff --git a/core/services/nodes/distributed_store_test.go b/core/services/nodes/distributed_store_test.go index 9b6e4ccc9464..08362fbfddca 100644 --- a/core/services/nodes/distributed_store_test.go +++ b/core/services/nodes/distributed_store_test.go @@ -2,11 +2,13 @@ package nodes import ( "context" + "errors" "fmt" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/LocalAI/pkg/model" ) @@ -48,15 +50,17 @@ var _ ModelLookup = (*fakeModelLookup)(nil) var _ = Describe("DistributedModelStore", func() { var ( - local *model.InMemoryModelStore - lookup *fakeModelLookup - store *DistributedModelStore + local *model.InMemoryModelStore + lookup *fakeModelLookup + clients *fakeBackendClientFactory + store *DistributedModelStore ) BeforeEach(func() { local = model.NewInMemoryModelStore() lookup = newFakeModelLookup() - store = NewDistributedModelStore(local, lookup) + clients = newFakeBackendClientFactory() + store = NewDistributedModelStore(local, lookup, clients) }) Describe("Get", func() { @@ -113,6 +117,63 @@ var _ = Describe("DistributedModelStore", func() { Expect(visited).To(HaveLen(2)) }) + It("gives every remote model a client that reaches the worker through its node", func() { + // The second construction path, closed. A model built with a nil + // client makes pkg/model.Model.GRPC dial its raw address with + // gRPC's own dialer the first time anything touches it, which + // bypasses the worker's tunnel completely. It is reached in + // production: ShutdownModel calls Free on it and the backend + // monitor calls Status. + dbNode := &BackendNode{ID: "node-2", Address: "10.0.0.3:50051"} + lookup.nodes["node-2"] = dbNode + lookup.allModels = []NodeModel{{NodeID: "node-2", ModelName: "remote-model"}} + + var got *model.Model + store.Range(func(id string, m *model.Model) bool { + if id == "remote-model" { + got = m + } + return true + }) + Expect(got).ToNot(BeNil()) + // The client is the factory's, so GRPC() returns it rather than + // building one by dialling. Asked for by NODE, not by address. + Expect(got.GRPC(false, nil)).To(BeIdenticalTo(grpc.Backend(clients.defaultClient))) + Expect(clients.nodesSeen()).To(ContainElement("node-2")) + }) + + It("refuses to list a remote model it has no way to reach", func() { + // Loudly, not by falling back. A model handed back here with a + // direct-dialling client works on a single-host developer setup and + // fails against every worker with no inbound port, which is the + // worst way for this defect to behave. + clients.refuseForNode = errors.New("no tunnel for you") + dbNode := &BackendNode{ID: "node-2", Address: "10.0.0.3:50051"} + lookup.nodes["node-2"] = dbNode + lookup.allModels = []NodeModel{{NodeID: "node-2", ModelName: "remote-model"}} + + visited := map[string]bool{} + store.Range(func(id string, _ *model.Model) bool { + visited[id] = true + return true + }) + Expect(visited).ToNot(HaveKey("remote-model")) + }) + + It("refuses when no client factory was wired at all", func() { + bare := NewDistributedModelStore(local, lookup, nil) + dbNode := &BackendNode{ID: "node-2", Address: "10.0.0.3:50051"} + lookup.nodes["node-2"] = dbNode + lookup.allModels = []NodeModel{{NodeID: "node-2", ModelName: "remote-model"}} + + visited := map[string]bool{} + bare.Range(func(id string, _ *model.Model) bool { + visited[id] = true + return true + }) + Expect(visited).ToNot(HaveKey("remote-model")) + }) + It("handles DB list error gracefully", func() { localModel := model.NewModel("model-x", "10.0.0.1:50051", nil) local.Set("model-x", localModel) diff --git a/core/services/nodes/file_stager_http.go b/core/services/nodes/file_stager_http.go index 79047aad6612..bd286e5d3426 100644 --- a/core/services/nodes/file_stager_http.go +++ b/core/services/nodes/file_stager_http.go @@ -28,9 +28,21 @@ import ( // Files are transferred between the frontend and backend nodes over a small // HTTP server running alongside the gRPC backend process. type HTTPFileStager struct { - httpAddrFor func(nodeID string) (string, error) - token string - client *http.Client + httpAddrFor func(nodeID string) (string, error) + token string + // dialFor supplies the transport for one worker. It is per node because a + // worker is reached over ITS OWN tunnel, and an http.Transport carries one + // DialContext: one shared transport could only ever reach one worker. + // + // nil means no tunnel dialer is wired, and every request is then refused + // rather than sent to the worker's advertised address; see + // ErrNoWorkerDialer for why that is not a fallback. + dialFor WorkerNetDialerFor + // clients caches one *http.Client per node. Caching is what keeps the + // connection pool: a client built per request would open a fresh tunnel + // stream for every chunk of a multi-gigabyte upload. + clientsMu sync.Mutex + clients map[string]*http.Client responseTimeout time.Duration // timeout waiting for server response after upload maxRetries int // number of retry attempts for transient failures } @@ -38,7 +50,8 @@ type HTTPFileStager struct { // NewHTTPFileStager creates a new HTTP file stager. // httpAddrFor should return the HTTP address (host:port) for the given node ID. // token is the registration token used for authentication. -func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token string) *HTTPFileStager { +// dialFor supplies the per-node transport; see the dialFor field. +func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token string, dialFor WorkerNetDialerFor) *HTTPFileStager { responseTimeout := 30 * time.Minute if v := os.Getenv("LOCALAI_FILE_TRANSFER_TIMEOUT"); v != "" { if d, err := time.ParseDuration(v); err == nil { @@ -53,11 +66,49 @@ func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token st } } + return &HTTPFileStager{ + httpAddrFor: httpAddrFor, + token: token, + dialFor: dialFor, + clients: map[string]*http.Client{}, + responseTimeout: responseTimeout, + maxRetries: maxRetries, + } +} + +// clientFor returns the HTTP client that reaches one worker, building it on +// first use. +// +// Every setting below is carried over unchanged from the single shared client +// this replaced, except DialContext, which now opens a stream on that worker's +// tunnel instead of connecting to its advertised address. HTTP/2 stays off for +// the reason it always was: its flow control stalls large uploads. +// +// What the tunnel dial does NOT carry over is the net.Dialer's own 30s connect +// timeout and 15s keepalive, because neither has anything left to act on: there +// is no TCP connect to time out, and liveness on the link is the yamux +// session's keepalive rather than the socket's. What still bounds a request is +// the context the caller passes. +// +// No client.Timeout is set, and that is deliberate: for large uploads +// http.Client.Timeout covers the whole request including the body, and firing +// mid-write closes the connection and shows up server-side as "connection reset +// by peer". The upload loop's own resume budget bounds the transfer instead. +func (h *HTTPFileStager) clientFor(nodeID string) (*http.Client, error) { + if h.dialFor == nil { + return nil, fmt.Errorf("staging files to node %s: %w", nodeID, ErrNoWorkerDialer) + } + h.clientsMu.Lock() + defer h.clientsMu.Unlock() + if c, ok := h.clients[nodeID]; ok { + return c, nil + } + dial := h.dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("staging files to node %s: %w", nodeID, ErrNoWorkerDialer) + } transport := &http.Transport{ - DialContext: (&net.Dialer{ - Timeout: 30 * time.Second, - KeepAlive: 15 * time.Second, // aggressive keepalive for LAN transfers - }).DialContext, + DialContext: dial, ForceAttemptHTTP2: false, // HTTP/2 flow control can stall large uploads MaxIdleConns: 10, IdleConnTimeout: 90 * time.Second, @@ -66,19 +117,9 @@ func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token st WriteBufferSize: 256 << 10, // 256 KB ReadBufferSize: 256 << 10, // 256 KB } - - return &HTTPFileStager{ - httpAddrFor: httpAddrFor, - token: token, - // No Timeout set — for large uploads, http.Client.Timeout covers the - // entire request lifecycle including the body upload. If it fires - // mid-write, Go closes the connection causing "connection reset by peer" - // on the server. Instead we use ResponseHeaderTimeout on the transport - // to cover only the wait-for-server-response phase. - client: httpclient.New(httpclient.WithTransport(transport)), - responseTimeout: responseTimeout, - maxRetries: maxRetries, - } + c := httpclient.New(httpclient.WithTransport(transport)) + h.clients[nodeID] = c + return c, nil } func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) { @@ -88,9 +129,13 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke if err != nil { return "", fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err) } + 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 := h.probeExisting(ctx, addr, localPath, key); ok { + if remotePath, ok := h.probeExisting(ctx, client, addr, localPath, key); ok { xlog.Info("Upload skipped (file already exists with matching hash)", "node", nodeID, "key", key, "remotePath", remotePath) return remotePath, nil } @@ -148,9 +193,9 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke // matching ours unlocks resume from the reported size; any other // outcome (missing file, hash mismatch, partial-of-different-file) // resets to 0 and uploads the entire file. - startOffset := h.resumeOffset(resumeCtx, addr, key, localHash, fileSize) + startOffset := h.resumeOffset(resumeCtx, client, addr, key, localHash, fileSize) - result, err := h.doUpload(ctx, resumeCtx, addr, nodeID, localPath, key, url, fileSize, startOffset, localHash) + result, err := h.doUpload(ctx, resumeCtx, client, addr, nodeID, localPath, key, url, fileSize, startOffset, localHash) if err == nil { if attempt > 1 { xlog.Info("File upload succeeded after retry", "node", nodeID, "file", filepath.Base(localPath), "attempt", attempt) @@ -237,7 +282,7 @@ func nextBackoff(attempt int) time.Duration { // different target hash). It returns the server-reported size when the // server's X-Target-SHA256 matches our expected final hash AND the size is // strictly less than the local file size. -func (h *HTTPFileStager) resumeOffset(ctx context.Context, addr, key, localHash string, fileSize int64) int64 { +func (h *HTTPFileStager) resumeOffset(ctx context.Context, client *http.Client, addr, key, localHash string, fileSize int64) int64 { if localHash == "" || fileSize <= 0 { return 0 } @@ -249,7 +294,7 @@ func (h *HTTPFileStager) resumeOffset(ctx context.Context, addr, key, localHash if h.token != "" { req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.client.Do(req) + resp, err := client.Do(req) if err != nil { return 0 } @@ -282,7 +327,7 @@ func (h *HTTPFileStager) resumeOffset(ctx context.Context, addr, key, localHash // the bytes from startOffset to fileSize-1. The outerCtx is the long-lived // resume budget; reqCtx is what's bound to the request (currently the same as // the parent ctx, since http.Client doesn't expose a per-request timeout). -func (h *HTTPFileStager) doUpload(ctx, outerCtx context.Context, addr, nodeID, localPath, key, url string, fileSize, startOffset int64, expectedHash string) (string, error) { +func (h *HTTPFileStager) doUpload(ctx, outerCtx context.Context, client *http.Client, addr, nodeID, localPath, key, url string, fileSize, startOffset int64, expectedHash string) (string, error) { if startOffset < 0 || startOffset > fileSize { startOffset = 0 } @@ -337,7 +382,7 @@ func (h *HTTPFileStager) doUpload(ctx, outerCtx context.Context, addr, nodeID, l req.Header.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", startOffset, fileSize-1, fileSize)) } - resp, err := h.client.Do(req) + resp, err := client.Do(req) if err != nil { xlog.Error("File upload failed", "node", nodeID, "file", filepath.Base(localPath), "size", humanFileSize(fileSize), "offset", startOffset, "error", err) @@ -441,7 +486,7 @@ func isTransientError(err error) bool { // file with a matching SHA-256 hash. Returns the remote path and true if the // upload can be skipped. Any errors (including 405 from older servers) silently // fall through so the caller proceeds with a normal PUT. -func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key string) (string, bool) { +func (h *HTTPFileStager) probeExisting(ctx context.Context, client *http.Client, addr, localPath, key string) (string, bool) { url := fmt.Sprintf("http://%s/v1/files/%s", addr, key) req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil) @@ -452,7 +497,7 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.client.Do(req) + resp, err := client.Do(req) if err != nil { return "", false } @@ -664,6 +709,10 @@ func (h *HTTPFileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, loca if err != nil { return fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err) } + client, err := h.clientFor(nodeID) + if err != nil { + return err + } if err := os.MkdirAll(filepath.Dir(localDst), 0750); err != nil { return fmt.Errorf("creating directory for %s: %w", localDst, err) @@ -680,7 +729,7 @@ func (h *HTTPFileStager) FetchRemoteByKey(ctx context.Context, nodeID, key, loca req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.client.Do(req) + resp, err := client.Do(req) if err != nil { return fmt.Errorf("downloading from node %s: %w", nodeID, err) } @@ -726,6 +775,10 @@ func (h *HTTPFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (st if err != nil { return "", fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err) } + client, err := h.clientFor(nodeID) + if err != nil { + return "", err + } url := fmt.Sprintf("http://%s/v1/files/temp", addr) req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil) @@ -736,7 +789,7 @@ func (h *HTTPFileStager) AllocRemoteTemp(ctx context.Context, nodeID string) (st req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.client.Do(req) + resp, err := client.Do(req) if err != nil { return "", fmt.Errorf("allocating temp file on node %s: %w", nodeID, err) } @@ -767,6 +820,10 @@ func (h *HTTPFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix st if err != nil { return nil, fmt.Errorf("resolving HTTP address for node %s: %w", nodeID, err) } + client, err := h.clientFor(nodeID) + if err != nil { + return nil, err + } url := fmt.Sprintf("http://%s/v1/files-list/%s", addr, keyPrefix) req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) @@ -777,7 +834,7 @@ func (h *HTTPFileStager) ListRemoteDir(ctx context.Context, nodeID, keyPrefix st req.Header.Set("Authorization", "Bearer "+h.token) } - resp, err := h.client.Do(req) + resp, err := client.Do(req) if err != nil { return nil, fmt.Errorf("listing dir on node %s: %w", nodeID, err) } diff --git a/core/services/nodes/file_stager_verify_deadline_test.go b/core/services/nodes/file_stager_verify_deadline_test.go index 0827bbecdc56..0239b832e7f8 100644 --- a/core/services/nodes/file_stager_verify_deadline_test.go +++ b/core/services/nodes/file_stager_verify_deadline_test.go @@ -65,7 +65,7 @@ var _ = Describe("staging verify phase and the cold-load stall window", func() { return "", err } return u.Host, nil - }, "") + }, "", directNetDialerFor) } It("survives a run of verified-and-skipped shards that upload no bytes at all", func() { diff --git a/core/services/nodes/file_transfer_server_test.go b/core/services/nodes/file_transfer_server_test.go index 78afb293b777..383d04b27bb9 100644 --- a/core/services/nodes/file_transfer_server_test.go +++ b/core/services/nodes/file_transfer_server_test.go @@ -21,6 +21,50 @@ import ( . "github.com/onsi/gomega" ) +// directNetDialerFor is the dial function these specs give the stager. +// +// The stager exists to reach a worker over that worker's TUNNEL, and it refuses +// to reach one at all without a dialer. These specs are about the HTTP protocol +// between the stager and the file-transfer server, and they run that server on +// loopback, so a plain TCP dial is what stands in for the tunnel here. Nothing +// in production supplies this: see the wiring in core/application. +func directNetDialerFor(_ string) func(ctx context.Context, network, addr string) (net.Conn, error) { + var d net.Dialer + return d.DialContext +} + +var _ = Describe("The HTTP file stager without a worker dialer", func() { + // Every request refused, none sent. Staging reaches a worker over that + // worker's tunnel, and a stager that fell back to connecting to the + // registered address would move gigabytes over a path that exists only + // while workers still listen on one. + newBare := func() *HTTPFileStager { + return NewHTTPFileStager(func(string) (string, error) { return "127.0.0.1:1", nil }, "tok", nil) + } + + It("refuses to upload", func() { + local := filepath.Join(GinkgoT().TempDir(), "f.bin") + Expect(os.WriteFile(local, []byte("payload"), 0o600)).To(Succeed()) + _, err := newBare().EnsureRemote(context.Background(), "node-1", local, "f.bin") + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + + It("refuses to download", func() { + dst := filepath.Join(GinkgoT().TempDir(), "out.bin") + Expect(newBare().FetchRemoteByKey(context.Background(), "node-1", "f.bin", dst)).To(MatchError(ErrNoWorkerDialer)) + }) + + It("refuses to allocate a remote temp file", func() { + _, err := newBare().AllocRemoteTemp(context.Background(), "node-1") + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + + It("refuses to list a remote directory", func() { + _, err := newBare().ListRemoteDir(context.Background(), "node-1", "models/") + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) +}) + var _ = Describe("FileTransferServer", func() { setupTestServer := func(token string, maxUploadSize int64) (*httptest.Server, string, string, string) { stagingDir := GinkgoT().TempDir() @@ -459,7 +503,7 @@ var _ = Describe("FileTransferServer", func() { addr := strings.TrimPrefix(ts.URL, "http://") stager := NewHTTPFileStager(func(nodeID string) (string, error) { return addr, nil - }, "tok") + }, "tok", directNetDialerFor) remotePath, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "present.bin") Expect(err).ToNot(HaveOccurred()) @@ -488,7 +532,7 @@ var _ = Describe("FileTransferServer", func() { addr := strings.TrimPrefix(ts.URL, "http://") stager := NewHTTPFileStager(func(nodeID string) (string, error) { return addr, nil - }, "tok") + }, "tok", directNetDialerFor) remotePath, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "changed.bin") Expect(err).ToNot(HaveOccurred()) @@ -517,7 +561,7 @@ var _ = Describe("FileTransferServer", func() { addr := strings.TrimPrefix(ts.URL, "http://") stager := NewHTTPFileStager(func(nodeID string) (string, error) { return addr, nil - }, "tok") + }, "tok", directNetDialerFor) remotePath, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "new.bin") Expect(err).ToNot(HaveOccurred()) @@ -553,7 +597,7 @@ var _ = Describe("FileTransferServer", func() { addr := strings.TrimPrefix(ts.URL, "http://") stager := NewHTTPFileStager(func(nodeID string) (string, error) { return addr, nil - }, "") + }, "", directNetDialerFor) remotePath, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "compat.bin") Expect(err).ToNot(HaveOccurred()) @@ -770,7 +814,7 @@ var _ = Describe("FileTransferServer", func() { addr := strings.TrimPrefix(ts.URL, "http://") stager := NewHTTPFileStager(func(nodeID string) (string, error) { return addr, nil - }, "tok") + }, "tok", directNetDialerFor) remotePath, err := stager.EnsureRemote(context.Background(), "node-1", localPath, "resume.bin") Expect(err).ToNot(HaveOccurred()) @@ -868,7 +912,7 @@ var _ = Describe("FileTransferServer", func() { addr := strings.TrimPrefix(ts.URL, "http://") stager := NewHTTPFileStager(func(nodeID string) (string, error) { return addr, nil - }, "tok") + }, "tok", directNetDialerFor) ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() diff --git a/core/services/nodes/health.go b/core/services/nodes/health.go index ffe1cfa0e2e5..44d79e0d8ab3 100644 --- a/core/services/nodes/health.go +++ b/core/services/nodes/health.go @@ -48,7 +48,11 @@ type HealthMonitor struct { // NewHealthMonitor creates a new HealthMonitor. // If db is non-nil (PostgreSQL), an advisory lock is used so that only one // frontend instance runs health checks at a time in distributed mode. -// If clientFactory is nil, a default factory using the given authToken is used. +// clientFactory is what reaches a worker's backends, over that worker's tunnel. +// Omitting it (or passing nil) leaves the monitor with a factory that refuses +// every request, so per-model probes are skipped and logged rather than counted +// as misses; authToken is then only the credential a working factory would have +// carried. Production always passes one. func NewHealthMonitor(registry NodeHealthStore, db *gorm.DB, checkInterval, staleThreshold time.Duration, authToken string, perModelHealthCheck bool, clientFactory ...BackendClientFactory) *HealthMonitor { checkInterval = cmp.Or(checkInterval, 15*time.Second) staleThreshold = cmp.Or(staleThreshold, 60*time.Second) @@ -184,7 +188,18 @@ func (hm *HealthMonitor) doCheckAll(ctx context.Context) { if m.Address == "" || m.Address == node.Address { continue } - mClient := hm.clientFactory.NewClient(m.Address, false) + // Through the node's tunnel, never a direct dial to m.Address: + // that address is a port inside the worker. A worker this + // replica cannot reach is not evidence that its backend died, + // so the miss counter is left alone and the row survives; + // counting it as a miss would reap live models across the whole + // fleet the moment the tunnel wiring was wrong. + mClient, err := hm.clientFactory.NewClientForNode(node.ID, m.Address, false) + if err != nil { + xlog.Error("Skipping model health probe: no way to reach the worker", + "node", node.ID, "model", m.ModelName, "replica", m.ReplicaIndex, "error", err) + continue + } mCheckCtx, mCancel := context.WithTimeout(ctx, 5*time.Second) ok, _ := mClient.HealthCheck(mCheckCtx) mCancel() diff --git a/core/services/nodes/health_mock_test.go b/core/services/nodes/health_mock_test.go index c52712dab5ff..281b876a7e4c 100644 --- a/core/services/nodes/health_mock_test.go +++ b/core/services/nodes/health_mock_test.go @@ -300,6 +300,11 @@ type fakeBackendClientFactory struct { clients map[string]*fakeBackendClient // default client returned when address not in clients map defaultClient *fakeBackendClient + // forNode records every node id NewClientForNode was asked for. + forNode []string + // refuseForNode makes NewClientForNode fail, standing in for a deployment + // with no way to reach the worker. Set before the code under test runs. + refuseForNode error } func newFakeBackendClientFactory() *fakeBackendClientFactory { @@ -324,6 +329,24 @@ func (f *fakeBackendClientFactory) NewClient(address string, _ bool) grpc.Backen return f.defaultClient } +// nodesSeen records the node ids the code under test asked for, so a spec can +// assert a caller reached a worker through its NODE rather than by address. +func (f *fakeBackendClientFactory) nodesSeen() []string { + f.mu.Lock() + defer f.mu.Unlock() + return append([]string(nil), f.forNode...) +} + +func (f *fakeBackendClientFactory) NewClientForNode(nodeID, address string, parallel bool) (grpc.Backend, error) { + if f.refuseForNode != nil { + return nil, f.refuseForNode + } + f.mu.Lock() + f.forNode = append(f.forNode, nodeID) + f.mu.Unlock() + return f.NewClient(address, parallel), nil +} + // helper to make a BackendNode with given properties func makeTestNode(id, name, address string, status string, lastHeartbeat time.Time) *BackendNode { return &BackendNode{ diff --git a/core/services/nodes/health_test.go b/core/services/nodes/health_test.go index c78ccfffe0d4..6f624b0cacb5 100644 --- a/core/services/nodes/health_test.go +++ b/core/services/nodes/health_test.go @@ -285,6 +285,46 @@ var _ = Describe("HealthMonitor (mock-based)", func() { Expect(store.getCalls()).NotTo(ContainElement(ContainSubstring("MarkUnhealthy"))) }) + It("probes a model through its NODE, never by dialling the stored address", func() { + // The address on a NodeModel row is a port inside the worker. This + // frontend reaches it over the worker's tunnel, so the node has to + // be part of every probe; a probe built from the address alone is + // the direct dial the tunnel replaces. + store := newFakeNodeHealthStore() + factory := newFakeBackendClientFactory() + hm := newTestHealthMonitor(store, factory, true, staleThreshold) + hm.perModelHealthCheck = true + + node := makeTestNode("node-tun", "tun-worker", "10.0.0.20:50051", StatusHealthy, freshTime()) + store.addNode(node) + store.addNodeModel("node-tun", NodeModel{NodeID: "node-tun", ModelName: "m", Address: "10.0.0.20:50053"}) + + hm.doCheckAll(context.Background()) + Expect(factory.nodesSeen()).To(ContainElement("node-tun")) + }) + + It("leaves a model row alone when it cannot reach the worker at all", func() { + // Not a miss. A frontend with no way to reach a worker has learned + // nothing about that worker's backends, and counting it as a failed + // probe would reap every model in the fleet the moment the tunnel + // wiring broke. + store := newFakeNodeHealthStore() + factory := newFakeBackendClientFactory() + factory.refuseForNode = fmt.Errorf("no tunnel for you") + hm := newTestHealthMonitor(store, factory, true, staleThreshold) + hm.perModelHealthCheck = true + + node := makeTestNode("node-cut", "cut-worker", "10.0.0.21:50051", StatusHealthy, freshTime()) + store.addNode(node) + store.addNodeModel("node-cut", NodeModel{NodeID: "node-cut", ModelName: "m", Address: "10.0.0.21:50053"}) + + for i := 0; i < perModelMissThreshold+1; i++ { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).NotTo(ContainElement(ContainSubstring("RemoveNodeModel"))) + Expect(store.getNode("node-cut").Status).To(Equal(StatusHealthy)) + }) + It("preserves model row when an intermittent failure is followed by a success", func() { store := newFakeNodeHealthStore() factory := newFakeBackendClientFactory() diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index be4dbc25d916..1231c811f371 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -2,6 +2,9 @@ package nodes import ( "context" + "errors" + "fmt" + "net" "time" "github.com/mudler/LocalAI/core/services/messaging" @@ -137,20 +140,96 @@ type NodeManager interface { RemoveAllNodeModelReplicas(ctx context.Context, nodeID, modelName string) error } -// BackendClientFactory creates gRPC backend clients. +// WorkerDialerFor hands back the dial function for one worker's backend +// processes: the shape grpc.WithContextDialer wants, bound to a node. +// +// It is declared here rather than taken as a core/services/cluster type so that +// this package keeps no dependency on that one. cluster is a leaf and imports +// neither this package nor core/http; the wiring in core/application supplies +// (*cluster.WorkerDialer).GRPCDialerFor, which has exactly this shape. +type WorkerDialerFor func(nodeID string) func(ctx context.Context, addr string) (net.Conn, error) + +// WorkerNetDialerFor hands back the dial function for one worker's own HTTP +// server, in the shape http.Transport.DialContext and websocket.Dialer's +// NetDialContext want. (*cluster.WorkerDialer).DialerFor bound to the http tag +// has this shape. +type WorkerNetDialerFor func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) + +// ErrNoWorkerDialer reports that something tried to reach a worker without a +// way to reach it through the worker's tunnel. +// +// It is deliberately an ERROR and not a fallback to dialling the worker's +// advertised address. A worker that holds a tunnel need not listen on anything +// and may be behind NAT with no address to dial, so the fallback would work +// only where the tunnel was not needed: on a single-host developer setup, and +// nowhere the feature exists for. It is also not an absence error: nothing +// about it says the worker is gone. +var ErrNoWorkerDialer = errors.New("nodes: no worker tunnel dialer is configured, so this worker cannot be reached") + +// BackendClientFactory creates the gRPC clients this frontend uses to reach +// model backends running on worker nodes. +// +// There is ONE method, and that is the design rather than an omission. A +// direct-dial constructor alongside it would be reachable from every call site +// that has an address, which is all of them, and the whole of this change is +// that having an address is no longer enough to reach a backend. Callers that +// genuinely want a raw address call pkg/grpc directly and are visible as such. type BackendClientFactory interface { - NewClient(address string, parallel bool) grpc.Backend + // NewClientForNode reaches a backend process running on a WORKER, through + // that worker's tunnel. address names WHICH process on the worker; it is + // not somewhere this process connects to. + // + // It returns an error rather than a client that falls back to a direct + // dial, so that a deployment with no tunnel dialer fails where the mistake + // is instead of quietly reopening the bypass. + NewClientForNode(nodeID, address string, parallel bool) (grpc.Backend, error) } -// tokenClientFactory is the default BackendClientFactory that creates gRPC -// clients with an optional bearer token for distributed auth. +// tokenClientFactory is the BackendClientFactory for a deployment with no +// worker tunnel dialer, which is a misconfiguration rather than a mode. It +// refuses every request, loudly, and reaches no worker. +// +// It exists so that the components that take a factory have something to hold +// when none was wired, instead of a nil they would have to guard at every use. type tokenClientFactory struct { token string } -func (f *tokenClientFactory) NewClient(address string, parallel bool) grpc.Backend { - if f.token != "" { - return grpc.NewClientWithToken(address, parallel, nil, false, f.token) +// NewClientForNode refuses. See ErrNoWorkerDialer for why this is not a direct +// dial to address. The token this factory carries is the one a working dialer +// would have used, kept only so the misconfiguration is repairable by wiring a +// dialer rather than by also re-plumbing credentials. +func (f *tokenClientFactory) NewClientForNode(nodeID, address string, _ bool) (grpc.Backend, error) { + return nil, fmt.Errorf("reaching backend %q on node %q: %w", address, nodeID, ErrNoWorkerDialer) +} + +// tunnelClientFactory reaches a worker's backend processes through the worker's +// tunnel, and is what every distributed deployment uses. +type tunnelClientFactory struct { + token string + dialFor WorkerDialerFor +} + +// NewTunnelClientFactory returns the factory that reaches worker backends +// through dialFor. A nil dialFor is refused rather than degraded: this +// constructor exists to close the direct-dial bypass, and one that silently +// handed back a direct-dialling factory would reopen it for the whole process. +func NewTunnelClientFactory(token string, dialFor WorkerDialerFor) (BackendClientFactory, error) { + if dialFor == nil { + return nil, fmt.Errorf("building the worker backend client factory: %w", ErrNoWorkerDialer) + } + return &tunnelClientFactory{token: token, dialFor: dialFor}, nil +} + +func (f *tunnelClientFactory) NewClientForNode(nodeID, address string, parallel bool) (grpc.Backend, error) { + if nodeID == "" { + // Without a node there is no tunnel to pick, and the only thing left to + // do with the address would be to dial it. + return nil, fmt.Errorf("reaching backend %q: no node id: %w", address, ErrNoWorkerDialer) + } + dial := f.dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("reaching backend %q on node %q: %w", address, nodeID, ErrNoWorkerDialer) } - return grpc.NewClient(address, parallel, nil, false) + return grpc.NewClientWithDialer(address, parallel, nil, false, f.token, dial), nil } diff --git a/core/services/nodes/local_stub_invalidator_test.go b/core/services/nodes/local_stub_invalidator_test.go index 00ed820dc6c2..444fefb6a9b2 100644 --- a/core/services/nodes/local_stub_invalidator_test.go +++ b/core/services/nodes/local_stub_invalidator_test.go @@ -44,7 +44,7 @@ var _ = Describe("LocalStubInvalidator", func() { }) It("drops the local stub once the last replica of the model is gone", func() { - store := NewDistributedModelStore(local, registry) + store := NewDistributedModelStore(local, registry, newFakeBackendClientFactory()) Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "ghost-model", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed()) local.Set("ghost-model", model.NewModel("ghost-model", "10.0.0.1:12345", nil)) @@ -64,7 +64,7 @@ var _ = Describe("LocalStubInvalidator", func() { }) It("keeps the local stub while another replica still serves the model", func() { - store := NewDistributedModelStore(local, registry) + store := NewDistributedModelStore(local, registry, newFakeBackendClientFactory()) Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "shared-model", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed()) Expect(registry.SetNodeModel(context.Background(), nodeB.ID, "shared-model", 0, "loaded", "10.0.0.2:12345", 0)).To(Succeed()) local.Set("shared-model", model.NewModel("shared-model", "10.0.0.1:12345", nil)) @@ -92,7 +92,7 @@ var _ = Describe("LocalStubInvalidator", func() { }) It("drops the local stub when a whole node's replicas are removed", func() { - store := NewDistributedModelStore(local, registry) + store := NewDistributedModelStore(local, registry, newFakeBackendClientFactory()) Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "node-model", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed()) local.Set("node-model", model.NewModel("node-model", "10.0.0.1:12345", nil)) diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index 62cc73e1545e..d50b55cef716 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -11,7 +11,6 @@ import ( "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" - grpcclient "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/xlog" "github.com/nats-io/nats.go" "google.golang.org/grpc/codes" @@ -34,13 +33,28 @@ const ( // ProbeUnreachable: nothing is listening (connection refused), or the // backend answered and affirmatively reported itself unhealthy. ProbeUnreachable + // ProbeUnknown: the probe was never made, because this frontend has no way + // to reach the worker at all (no tunnel dialer wired, or none for this + // node). It is NOT ProbeUnreachable and must never be folded into it: + // unreachable is an observation about a backend and the reaper deletes rows + // on it, while this is a statement about THIS process and says nothing + // about the worker, which may be running the model perfectly well. + // + // It is appended rather than made the zero value on purpose. ProbeAlive is + // the zero value already, and renumbering the set would silently change the + // meaning of every stored or hard-coded outcome. + ProbeUnknown ) // ModelProber checks the state of a model's backend process. // Defaulted to a gRPC health probe but overridable for tests so we don't // need to stand up a real server. type ModelProber interface { - Probe(ctx context.Context, address string) ProbeOutcome + // Probe checks the backend at address on node nodeID. The node is needed + // as well as the address because address is a port INSIDE the worker, + // reached over the tunnel that worker holds, and there is no route to it + // that does not name the node. + Probe(ctx context.Context, nodeID, address string) ProbeOutcome } // NodeProcessLister asks a worker which model backend processes it currently @@ -60,11 +74,20 @@ type NodeProcessLister interface { // as death. const probeTimeout = 1 * time.Second -// grpcModelProber does a short HealthCheck on the model's stored gRPC address. -type grpcModelProber struct{ token string } +// grpcModelProber does a short HealthCheck on the model's stored gRPC address, +// through the tunnel of the node that address belongs to. +type grpcModelProber struct{ clients BackendClientFactory } -func (g grpcModelProber) Probe(ctx context.Context, address string) ProbeOutcome { - client := grpcclient.NewClientWithToken(address, false, nil, false, g.token) +func (g grpcModelProber) Probe(ctx context.Context, nodeID, address string) ProbeOutcome { + client, err := g.clients.NewClientForNode(nodeID, address, false) + if err != nil { + // Never ProbeUnreachable: the reaper deletes a row on that answer, and + // this frontend not being able to reach a worker is no evidence that + // the worker stopped running the model. + xlog.Error("Cannot probe a model: no way to reach the worker", + "node", nodeID, "address", address, "error", err) + return ProbeUnknown + } probeCtx, cancel := context.WithTimeout(ctx, probeTimeout) defer cancel() ok, err := client.HealthCheck(probeCtx) @@ -173,11 +196,17 @@ type ReplicaReconcilerOptions struct { // Adapter is the NATS sender used to retry pending backend ops. When nil, // the state-reconciler pending-drain pass is a no-op (single-node mode). Adapter *RemoteUnloaderAdapter - // RegistrationToken is used by the default gRPC prober when probing model - // addresses. Matches the worker's token so HealthCheck auth succeeds. + // RegistrationToken is the bearer token the default gRPC prober presents to + // a worker's backends. It matters only when ClientFactory is unset, since + // the factory carries its own; a prober built from the token alone can + // reach no worker at all and reports ProbeUnknown for every model. RegistrationToken string // Prober overrides the default gRPC health probe (used by tests). Prober ModelProber + // ClientFactory builds the gRPC clients the default prober uses. It is what + // carries the worker tunnel dialer; without it the default prober can reach + // no worker and says so on every probe. + 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. @@ -210,7 +239,14 @@ func NewReplicaReconciler(opts ReplicaReconcilerOptions) *ReplicaReconciler { } prober := opts.Prober if prober == nil { - prober = grpcModelProber{token: opts.RegistrationToken} + clients := opts.ClientFactory + if clients == nil { + // No tunnel dialer was wired. The prober then refuses every probe + // with ProbeUnknown rather than dialling addresses directly, which + // is loud in the log and leaves every row alone. + clients = &tokenClientFactory{token: opts.RegistrationToken} + } + prober = grpcModelProber{clients: clients} } pressureThreshold := opts.PressureThreshold if pressureThreshold == 0 { @@ -469,7 +505,15 @@ func (rc *ReplicaReconciler) probeLoadedModels(ctx context.Context) { return } seen[m.ID] = struct{}{} - switch rc.prober.Probe(ctx, m.Address) { + switch rc.prober.Probe(ctx, m.NodeID, m.Address) { + case ProbeUnknown: + // This frontend could not reach the worker to ask. The streak is + // left exactly as it was: neither cleared, which would forgive a + // backend that really is dead, nor advanced, which would reap every + // model in the fleet the moment the tunnel wiring broke. + xlog.Warn("Reconciler: could not probe a model, leaving its row alone", + "node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address) + continue case ProbeAlive: rc.clearProbeFailures(m.ID) // Bump updated_at so we don't probe this row again immediately. @@ -552,7 +596,7 @@ func (rc *ReplicaReconciler) sweepLeakedInFlight(ctx context.Context) { return } seen[m.ID] = struct{}{} - if rc.prober.Probe(ctx, m.Address) != ProbeAlive { + if rc.prober.Probe(ctx, m.NodeID, m.Address) != ProbeAlive { // Busy or unreachable. Busy means the counter may well be real; // unreachable is the reaper's business, not the sweeper's. rc.clearInFlightIdle(m.ID) diff --git a/core/services/nodes/reconciler_busy_probe_test.go b/core/services/nodes/reconciler_busy_probe_test.go index 9cc7b07f9923..c2a6b076836c 100644 --- a/core/services/nodes/reconciler_busy_probe_test.go +++ b/core/services/nodes/reconciler_busy_probe_test.go @@ -89,6 +89,48 @@ var _ = Describe("ReplicaReconciler — probe reaper vs busy backends", func() { "a backend that accepted the connection but was mid-request must never be reaped") }) + It("never reaps a replica it could not probe at all", func() { + // ProbeUnknown is this FRONTEND saying it has no way to reach the + // worker, which is nothing at all about the backend. Folding it into + // ProbeUnreachable would empty the whole node_models table the moment + // the tunnel wiring was wrong, and the models would still be running. + seed("unknown-1", 0) + prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeUnknown}} + rc := newReconciler(prober) + + for range probeFailuresBeforeReap * 3 { + rc.probeLoadedModels(context.Background()) + makeStale("unknown-1") + } + + var after NodeModel + Expect(db.First(&after, "id = ?", "unknown-1").Error).To(Succeed(), + "a replica this frontend could not reach must never be reaped") + }) + + It("does not let an unprobeable pass forgive a real failure streak", func() { + // The other half of the same rule. Clearing the streak on an outcome + // that observed nothing would let a flapping tunnel keep a genuinely + // dead backend in the table forever. + seed("mixed-1", 0) + prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeUnreachable}} + rc := newReconciler(prober) + + for i := 1; i < probeFailuresBeforeReap; i++ { + rc.probeLoadedModels(context.Background()) + makeStale("mixed-1") + } + prober.outcomes[addr] = ProbeUnknown + rc.probeLoadedModels(context.Background()) + makeStale("mixed-1") + + prober.outcomes[addr] = ProbeUnreachable + rc.probeLoadedModels(context.Background()) + var after NodeModel + Expect(db.First(&after, "id = ?", "mixed-1").Error).To(MatchError(gorm.ErrRecordNotFound), + "an unprobeable pass must leave the streak untouched, not reset it") + }) + It("reaps an unreachable replica even when in_flight leaked high", func() { // in_flight has no decrement guarantee: a frontend that dies mid-request // leaves the increment behind forever. Gating the reaper on it would diff --git a/core/services/nodes/reconciler_test.go b/core/services/nodes/reconciler_test.go index 049fb94418de..1ed0b47d8ad7 100644 --- a/core/services/nodes/reconciler_test.go +++ b/core/services/nodes/reconciler_test.go @@ -740,7 +740,7 @@ type fakeProber struct { calls int } -func (f *fakeProber) Probe(_ context.Context, address string) ProbeOutcome { +func (f *fakeProber) Probe(_ context.Context, _, address string) ProbeOutcome { f.calls++ if f.outcomes == nil { return ProbeUnreachable diff --git a/core/services/nodes/revision_eligibility_test.go b/core/services/nodes/revision_eligibility_test.go index 96ea5010b704..fd01c5b7c9cc 100644 --- a/core/services/nodes/revision_eligibility_test.go +++ b/core/services/nodes/revision_eligibility_test.go @@ -264,7 +264,7 @@ var _ = Describe("revision eligibility consumers", func() { type recordingEligibilityProber struct{ addresses []string } -func (p *recordingEligibilityProber) Probe(_ context.Context, address string) ProbeOutcome { +func (p *recordingEligibilityProber) Probe(_ context.Context, _, address string) ProbeOutcome { p.addresses = append(p.addresses, address) return ProbeAlive } diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index d094d0f2076e..b9731f4a305a 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -393,7 +393,10 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking } } - client := r.buildClientForAddr(node, backendAddr, parallel) + client, err := r.buildClientForAddr(node, backendAddr, parallel) + if err != nil { + return nil, fmt.Errorf("building a client for model %q on node %q: %w", modelName, node.ID, err) + } // Load the model on the remote node if loadOpts != nil { @@ -721,7 +724,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) { + alive, probed := r.probeHealth(ctx, node, modelAddr) + if !probed { + // Nothing was asked, so nothing was learned. The row is left exactly + // as it was: removing it would reclaim a model that is loaded and + // healthy on a worker this frontend merely cannot reach right now. The + // reservation is released, and the cold path below reports the wiring + // fault with the detail a caller needs. + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release a reservation for an unreachable worker", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + return nil + } + if !alive { // 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", @@ -753,7 +769,20 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route // call finishes, so in-flight returns to 0 when idle. r.registry.TouchNodeModel(ctx, node.ID, att.trackingKey, replicaIdx) r.observePrefix(att.trackingKey, att.observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) - grpcClient := r.buildClientForAddr(node, modelAddr, att.parallel) + grpcClient, err := r.buildClientForAddr(node, modelAddr, att.parallel) + if err != nil { + // The probe above builds a client for the same node and would have + // reported !probed, so reaching here means the dialer stopped being + // able to serve this node between the two. Handled the same way and for + // the same reason: release the reservation, leave the row alone. + if relErr := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); relErr != nil { + xlog.Warn("Failed to release a reservation for an unreachable worker", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", relErr) + } + xlog.Error("Cannot build a client for a loaded model: no way to reach the worker", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + return nil + } tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, att.trackingKey, replicaIdx) return r.newRouteResult(node, att.trackingKey, replicaIdx, grpcClient, tracked) } @@ -1343,14 +1372,25 @@ func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNod } } -func (r *SmartRouter) buildClientForAddr(node *BackendNode, addr string, parallel bool) grpc.Backend { - client := r.clientFactory.NewClient(addr, parallel) +// buildClientForAddr builds the gRPC client for a backend process running on a +// worker node. +// +// addr is a port INSIDE the worker, reached over the tunnel that worker holds; +// connecting to it from here would only work for a worker that still listens on +// a routable address. The factory offers no way to do that, and an error is +// returned rather than a direct-dialling client for the reason +// ErrNoWorkerDialer gives. +func (r *SmartRouter) buildClientForAddr(node *BackendNode, addr string, parallel bool) (grpc.Backend, error) { + client, err := r.clientFactory.NewClientForNode(node.ID, addr, parallel) + if err != nil { + return nil, err + } // Wrap with file staging if configured if r.fileStager != nil { - return NewFileStagingClient(client, r.fileStager, node.ID) + return NewFileStagingClient(client, r.fileStager, node.ID), nil } - return client + return client, nil } // stageModelFiles uploads model files to the backend node via the FileStager. @@ -1885,6 +1925,13 @@ func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir // via a gRPC health check with a 2-second timeout. The client is closed after // the check. // +// TWO results, not one. alive is what the backend said; probed is whether it +// was asked at all. They are separate because the caller REAPS on a dead probe, +// and a frontend that cannot reach a worker has observed nothing about that +// worker's backends: folding the two would delete every replica row in the +// deployment the moment the tunnel wiring was wrong, while the models carried +// on running. +// // The result is memoized in r.probeCache for probeCacheTTL. With per-request // routing every inference call lands here, and unbounded re-probing can stall // behind a busy backend that serializes HealthCheck against active Predict. @@ -1892,16 +1939,27 @@ 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 client is built OUTSIDE the memoized closure, which is what keeps an +// unreachable worker out of the cache entirely: DoOrCached only ever sees a +// real answer. Building it costs a struct and no I/O, since the gRPC client +// dials lazily on its first call. +func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr string) (alive, probed bool) { + client, err := r.buildClientForAddr(node, addr, false) + if err != nil { + xlog.Error("Cannot probe a model backend: no way to reach the worker", + "node", node.ID, "address", addr, "error", err) + return false, false + } + defer closeClient(client) + key := node.ID + "|" + addr return r.probeCache.DoOrCached(key, func() bool { - 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 - }) + }), true } // closeClient closes a gRPC backend client if it implements io.Closer. diff --git a/core/services/nodes/router_load_budget_test.go b/core/services/nodes/router_load_budget_test.go index 921b19905d4a..f0f41f78c235 100644 --- a/core/services/nodes/router_load_budget_test.go +++ b/core/services/nodes/router_load_budget_test.go @@ -86,6 +86,10 @@ type holdClientFactory struct{ client *holdBackend } func (f *holdClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } +func (f *holdClientFactory) NewClientForNode(_, address string, parallel bool) (grpc.Backend, error) { + return f.NewClient(address, parallel), nil +} + var _ = Describe("size-derived remote LoadModel budget", func() { // Production, on an NVIDIA Jetson Thor worker: a 70 GB video checkpoint // (longcat-video-avatar-1.5) failed reproducibly after 953.5s with diff --git a/core/services/nodes/router_load_timeout_test.go b/core/services/nodes/router_load_timeout_test.go index 295dc35d81fd..7ee68ac123e6 100644 --- a/core/services/nodes/router_load_timeout_test.go +++ b/core/services/nodes/router_load_timeout_test.go @@ -53,6 +53,10 @@ type deadlineClientFactory struct{ client *deadlineBackend } func (f *deadlineClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } +func (f *deadlineClientFactory) NewClientForNode(_, address string, parallel bool) (grpc.Backend, error) { + return f.NewClient(address, parallel), nil +} + var _ = Describe("remote LoadModel deadline", func() { var ( reg *fakeModelRouter diff --git a/core/services/nodes/router_reap_load_test.go b/core/services/nodes/router_reap_load_test.go index 67376c06f535..b2619dc97f0a 100644 --- a/core/services/nodes/router_reap_load_test.go +++ b/core/services/nodes/router_reap_load_test.go @@ -45,6 +45,10 @@ type failingClientFactory struct{ client *failingLoadBackend } func (f *failingClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } +func (f *failingClientFactory) NewClientForNode(_, address string, parallel bool) (grpc.Backend, error) { + return f.NewClient(address, parallel), nil +} + // replicaSlotRouter pins the replica slot scheduleAndLoad allocates so a spec // can assert the reaped process key carries the real index, not a hardcoded 0. type replicaSlotRouter struct { diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 015cacb3040b..68cd376246e3 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -466,6 +466,10 @@ func (f *stubClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } +func (f *stubClientFactory) NewClientForNode(_, address string, parallel bool) (grpc.Backend, error) { + return f.NewClient(address, parallel), nil +} + // --------------------------------------------------------------------------- // Fake NodeCommandSender (unloader) // --------------------------------------------------------------------------- diff --git a/core/services/nodes/router_unreachable_worker_test.go b/core/services/nodes/router_unreachable_worker_test.go new file mode 100644 index 000000000000..943741396e11 --- /dev/null +++ b/core/services/nodes/router_unreachable_worker_test.go @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +// unreachableClientFactory cannot build a client for any node, standing in for +// a frontend whose worker tunnel dialer is missing or broken. +type unreachableClientFactory struct{} + +func (unreachableClientFactory) NewClientForNode(_, _ string, _ bool) (grpc.Backend, error) { + return nil, errors.New("no way to reach that worker") +} + +var _ = Describe("routing when the worker cannot be reached at all", func() { + // The catastrophe this phase exists to prevent, at the router. A frontend + // that cannot reach a worker has learned NOTHING about that worker's + // backends. Treating it as a failed health probe would reap the replica row + // for every model in the deployment while those models carried on running, + // and the reap is silent: the row is simply deleted and the model + // cold-loaded somewhere else. + loadedReg := func() *fakeModelRouter { + node := &BackendNode{ID: "X", Name: "node-x", Address: "10.0.0.1:50051"} + nm := &NodeModel{NodeID: "X", ModelName: "m", Address: "10.0.0.1:9001"} + return &fakeModelRouter{ + findAndLockNode: node, + findAndLockNM: nm, + loadedReplicaStatsByName: map[string][]ReplicaCandidate{"m": {{NodeID: "X", InFlight: 0}}}, + } + } + + It("never removes the replica row of a worker it merely cannot reach", func() { + reg := loadedReg() + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: &fakeUnloader{}, + ClientFactory: unreachableClientFactory{}, + }) + + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) + // The request cannot be served, which is right and loud. + Expect(err).To(HaveOccurred()) + // What must NOT have happened is the replica being reclaimed. + Expect(reg.removeCalls).To(BeEmpty(), + "a worker this frontend cannot reach must never have its loaded models reaped") + }) + + It("releases the routing reservation it took before giving up", func() { + // FindAndLockNodeWithModel increments in_flight as a reservation. A + // path that returns without releasing it leaves the replica looking + // permanently busy, which is how a warm replica stops being picked at + // all. + reg := loadedReg() + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: &fakeUnloader{}, + ClientFactory: unreachableClientFactory{}, + }) + + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) + Expect(err).To(HaveOccurred()) + Expect(reg.decrementCalls).To(ContainElement("X:m")) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 22f61ba66025..d73cdf6b6cb6 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -143,6 +143,37 @@ A stream naming a tag the worker does not serve, a target outside that port rang Reconnects use exponential backoff with jitter: the interval doubles from 500ms up to a ceiling of 30 seconds, and each wait is drawn between half of that interval and all of it, so no worker ever spins and a fleet that lost the same replica does not come back in lockstep. The interval returns to its floor only after a session that lasted at least 30 seconds. That last part is what stops a rolling frontend restart, where every dial succeeds and then dies moments later, from turning a fleet of workers into a retry storm against the first replica back up. A worker that is refused (`401`, `403`) keeps retrying on the same schedule rather than exiting: a re-registration or an admin approval fixes both without restarting it. +#### What the frontend sends through it + +Every connection the frontend makes to a worker now goes through that worker's tunnel. There are three, and all three are the same path underneath: + +| What | Protocol | Stream tag | +|------|----------|-----------| +| Inference, model load, health checks | gRPC to a backend process | `grpc` | +| Model file staging, backend-log listing | HTTP to the worker's own server | `http` | +| Live backend-log streaming | WebSocket to the same server | `http` | + +The address the frontend holds for a backend (the per-replica port a worker reports after an install) is still what identifies it, and it is still what appears in logs and errors. What it no longer is, is somewhere the frontend connects to: it travels inside the tunnel as the stream's target, and the worker decides what to do with it. + +A frontend with no way to reach a worker says so and fails. It does **not** fall back to connecting to the worker's advertised address. That fallback is what the tunnel exists to remove, and it is the kind of defect that works on a one-replica developer box and fails in production, so it is an error everywhere. The consequences are deliberately narrow: a model whose worker cannot be reached is not reaped, and its row is left alone, because a frontend that cannot reach a worker has learned nothing about whether that worker is still running the model. + +#### Reaching a worker another replica holds + +A worker's tunnel lands on exactly one replica, so with N replicas behind a load balancer roughly (N-1)/N of requests arrive somewhere else. Those requests are relayed: the replica that received the request looks up the owner in `node_connections`, **joined against the live `instances` rows**, opens a stream on its peer link to that owner, and the owner splices it onto the worker's tunnel. One hop, never two; a stale ownership row is answered with a routing refusal and the dialling replica resolves the owner again rather than being sent round a loop. + +The dialling replica states how much time its own client has left in the frame that opens the relayed stream, and the owner bounds its work by the smaller of that and its own 15s ceiling. Neither number can lengthen the other: a patient client cannot park the owning replica, and an impatient one cannot be kept waiting on a budget it did not ask for. + +Four outcomes are kept apart on purpose, because they call for different actions: + +| Outcome | What it means | What acts on it | +|---|---|---| +| No live owner | No replica holds this worker's tunnel | The worker is treated as absent; its models can be rescheduled | +| Not the owner | The routing was stale | Resolve the owner again | +| Peer unreachable | A replica exists and will not answer | Retry | +| The worker refused | The worker answered and said no | Report; the worker is connected | + +Only the first is absence. The others are never reported as it, and that is not a stylistic preference: a scheduler told that a connected worker has gone away reclaims every model it is running. + Set `LOCALAI_WORKER_TUNNEL=false` on a worker to turn the tunnel off and go back to the frontend dialling the worker's advertised addresses. ### The model load deadline scales with the checkpoint diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 93dde00991b7..e90bbb28fc5f 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -2,6 +2,7 @@ package grpc import ( "context" + "net" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "google.golang.org/grpc" @@ -29,7 +30,30 @@ 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 rather than by connecting to address. +// +// It is what distributed mode uses to reach a backend process on a worker: the +// worker holds one multiplexed tunnel to a frontend replica and listens on +// nothing, so address names which backend process the stream is for and the +// dialer decides how the stream gets there. A nil dialer is a programming +// error on this path rather than a fallback, because falling back to a direct +// dial would work in a single-replica test and fail in production; callers with +// no dialer call NewClientWithToken and mean it. +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 rather than through a checked assertion: + // an assertion that failed would silently hand back a client that dials + // the address directly, which is the exact bypass this constructor exists + // to close. + c := buildClient(address, parallel, wd, enableWatchDog, token) + c.dialer = dialer + return c +} + +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 a6f8947eba61..6131d3801bc7 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,13 @@ type Client struct { inFlight int parallel bool token string + // dialer replaces the transport gRPC would otherwise use to reach address. + // In distributed mode it is a stream on the worker's tunnel, so address + // stops being a socket to connect to and becomes the name of a backend + // process inside the worker; see core/services/cluster.WorkerDialer. nil + // keeps gRPC's own TCP dial, which is what every non-distributed caller + // wants. + dialer func(ctx context.Context, addr string) (net.Conn, error) sync.Mutex opMutex sync.Mutex wd WatchDog @@ -80,6 +88,12 @@ func (c *Client) dial() (*grpc.ClientConn, error) { if c.token != "" { opts = append(opts, grpc.WithPerRPCCredentials(bearerToken{token: c.token})) } + if c.dialer != nil { + // The address is still passed to grpc.NewClient because it is what + // names the target in every error message and in the authority header; + // what it no longer decides is where the bytes go. + opts = append(opts, grpc.WithContextDialer(c.dialer)) + } return grpc.NewClient(c.address, opts...) } diff --git a/tests/e2e/distributed/backend_logs_test.go b/tests/e2e/distributed/backend_logs_test.go index 82e8ac156401..a858629ef464 100644 --- a/tests/e2e/distributed/backend_logs_test.go +++ b/tests/e2e/distributed/backend_logs_test.go @@ -50,6 +50,17 @@ func waitForSingleLogSubscriber(logStore *model.BackendLogStore, modelID string) Should(Equal(1), "the WebSocket handler never subscribed to %q exactly once", modelID) } +// directWorkerDialerFor stands in for the worker tunnel in these specs. +// +// The log-proxy endpoints reach a worker over the tunnel that worker holds, and +// refuse to reach one without a dialer. These specs run the worker's HTTP +// server on loopback, so a plain TCP dial is the stand-in; production supplies +// the real one from core/application. +func directWorkerDialerFor(_ string) func(ctx context.Context, network, addr string) (net.Conn, error) { + var d net.Dialer + return d.DialContext +} + var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func() { Context("Worker HTTP log endpoints", func() { @@ -370,7 +381,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func // Create an Echo test server with the proxy endpoint e := echo.New() - e.GET("/api/nodes/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, token)) + e.GET("/api/nodes/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, token, directWorkerDialerFor)) req := httptest.NewRequest("GET", fmt.Sprintf("/api/nodes/%s/backend-logs", node.ID), nil) rec := httptest.NewRecorder() @@ -392,7 +403,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(registry.Register(context.Background(), node, true)).To(Succeed()) e := echo.New() - e.GET("/api/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsLinesEndpoint(registry, token)) + e.GET("/api/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsLinesEndpoint(registry, token, directWorkerDialerFor)) req := httptest.NewRequest("GET", fmt.Sprintf("/api/nodes/%s/backend-logs/remote-model", node.ID), nil) rec := httptest.NewRecorder() @@ -409,7 +420,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func It("should return 404 for unknown node ID", func() { e := echo.New() - e.GET("/api/nodes/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, token)) + e.GET("/api/nodes/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, token, directWorkerDialerFor)) req := httptest.NewRequest("GET", "/api/nodes/nonexistent-id/backend-logs", nil) rec := httptest.NewRecorder() @@ -453,7 +464,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func // Start Echo server with the WebSocket proxy route e := echo.New() - e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, token)) + e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, token, directWorkerDialerFor)) lis, err := net.Listen("tcp", "127.0.0.1:0") Expect(err).ToNot(HaveOccurred()) diff --git a/tests/e2e/distributed/distributed_full_flow_test.go b/tests/e2e/distributed/distributed_full_flow_test.go index ad7f2669aaf0..de26573790df 100644 --- a/tests/e2e/distributed/distributed_full_flow_test.go +++ b/tests/e2e/distributed/distributed_full_flow_test.go @@ -489,7 +489,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() return "", err } return n.HTTPAddress, nil - }, "") + }, "", directWorkerDialerFor) // Create SmartRouter with the HTTPFileStager router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) @@ -558,7 +558,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() return "", err } return n.HTTPAddress, nil - }, "") + }, "", directWorkerDialerFor) // Create SmartRouter with FileStager router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) @@ -616,7 +616,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() return "", err } return n.HTTPAddress, nil - }, "") + }, "", directWorkerDialerFor) // Test AllocRemoteTemp + FetchRemote directly (the output retrieval path) remoteTmpPath, err := stager.AllocRemoteTemp(ctx, node.ID) @@ -662,7 +662,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() return "", err } return n.HTTPAddress, nil - }, "") + }, "", directWorkerDialerFor) router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) @@ -881,7 +881,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() return "", err } return n.HTTPAddress, nil - }, "") + }, "", directWorkerDialerFor) // Create model files on the "frontend" frontendModelsDir := GinkgoT().TempDir() @@ -965,7 +965,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() return "", err } return n.HTTPAddress, nil - }, "") + }, "", directWorkerDialerFor) // Create model files: .onnx and .onnx.json in a temp "models" dir frontendModelsDir := GinkgoT().TempDir() diff --git a/tests/e2e/distributed/distributed_store_test.go b/tests/e2e/distributed/distributed_store_test.go index 679a75a02b28..d35a7e0439a8 100644 --- a/tests/e2e/distributed/distributed_store_test.go +++ b/tests/e2e/distributed/distributed_store_test.go @@ -2,6 +2,7 @@ package distributed_test import ( "context" + "net" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/pkg/model" @@ -14,6 +15,25 @@ import ( "gorm.io/gorm/logger" ) +// directBackendClients stands in for the worker tunnel in these specs. +// +// The store refuses to build a client for a remote model without a way to reach +// the worker, which is the point: a model built with no client dials its raw +// address on first use. These specs have no worker tunnel and no worker, so the +// dial is a plain TCP one; production supplies the real dialer from +// core/application. +func directBackendClients() nodes.BackendClientFactory { + GinkgoHelper() + clients, err := nodes.NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + var d net.Dialer + return func(ctx context.Context, addr string) (net.Conn, error) { + return d.DialContext(ctx, "tcp", addr) + } + }) + Expect(err).ToNot(HaveOccurred()) + return clients +} + var _ = Describe("DistributedModelStore", Label("Distributed"), func() { var ( infra *TestInfra @@ -36,7 +56,7 @@ var _ = Describe("DistributedModelStore", Label("Distributed"), func() { Expect(err).ToNot(HaveOccurred()) localStore = model.NewInMemoryModelStore() - dStore = nodes.NewDistributedModelStore(localStore, registry) + dStore = nodes.NewDistributedModelStore(localStore, registry, directBackendClients()) }) Context("Get", func() { diff --git a/tests/e2e/distributed/file_staging_test.go b/tests/e2e/distributed/file_staging_test.go index 55bd5663c6b4..e765d87c3b87 100644 --- a/tests/e2e/distributed/file_staging_test.go +++ b/tests/e2e/distributed/file_staging_test.go @@ -62,7 +62,7 @@ var _ = Describe("File Staging", Label("Distributed"), func() { It("should create HTTPFileStager with httpAddrFor function", func() { stager := nodes.NewHTTPFileStager(func(nodeID string) (string, error) { return "", fmt.Errorf("no such node: %s", nodeID) - }, "") + }, "", directWorkerDialerFor) Expect(stager).ToNot(BeNil()) // Should fail gracefully when node resolution fails diff --git a/tests/e2e/distributed/prefix_cache_routing_test.go b/tests/e2e/distributed/prefix_cache_routing_test.go index 9b1e3c117718..e899160ea2fd 100644 --- a/tests/e2e/distributed/prefix_cache_routing_test.go +++ b/tests/e2e/distributed/prefix_cache_routing_test.go @@ -51,6 +51,10 @@ func (f *prefixStubClientFactory) NewClient(_ string, _ bool) grpcPkg.Backend { return f.client } +func (f *prefixStubClientFactory) NewClientForNode(_, _ string, _ bool) (grpcPkg.Backend, error) { + return f.client, nil +} + var _ = Describe("Prefix-cache aware routing", Label("Distributed"), func() { const model = "model" From a8ac2af167cfa4ece4715ae2ffc3a28fcc056ff8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 16:09:26 +0000 Subject: [PATCH 32/42] fix(cluster): make "no route" a condition of its own, and let it out of the package Review round 1 on task 6. Five blocking findings, all with the same root: the conditions the dialer kept apart were erased one layer out, because every one of them arrived at core/services/nodes as a gRPC codes.Unavailable, which is also what a backend process that died produces. Four call sites acted on that by deleting a replica row, one of them after a single failed probe. The fifth condition is ErrNoRoute: this replica could not get a request to a worker's backend, and no claim at all about the worker. A worker's presence is its HEARTBEAT, which nodes owns; a route is a separate fact that cluster owns, and the two now differ. They differ in normal operation, not exotically: a worker that has not dialled its tunnel yet after a frontend-first upgrade is unroutable on every request while it heartbeats and serves. Two properties, both mutation-tested. Every failure to resolve or open a route carries ErrNoRoute, so a consumer has one check to make. No failure carries an absence sentinel: routeFailure is the single place that rule lives, and it keeps ErrNoConnection and ErrInstanceNotFound in the message and out of the unwrap chain, the guarantee unreachableError already made for peers. Everything else stays matchable, so ErrNotOwner and ErrPeerUnreachable are unchanged for anyone who can act on them. A worker's own refusal carries no umbrella, because a worker that answers has demonstrated it is there and that is the only real evidence on the path. Crossing the boundary needed a value, not a code. NewClientWithDialer wraps the dialer and records each outcome; LastDialError hands it back behind a narrow interface, and nodes.unroutable turns it into ErrWorkerUnroutable with the cluster sentinels still in the chain. A spec asserts a dial failing with ErrNoRoute plus ErrPeerUnreachable arrives matching all three and matching neither absence sentinel. The sweep found a fourth site the review had not named: pkg/model checkIsLoaded evicts a remote model on a connection error, and a tunnel dial failure is one. Four other reap sites were cleared with reasons - inflight and the worker authoritative pass reap only on semantic answers, scale-down is driven by last_used, abandoned loads decide on the node's heartbeat. Every fixed site also grew the opposite spec, so the new check cannot pass by never reaping. probeCache carries the reason through singleflight rather than a closed-over variable. A variable is only written by the goroutine that runs the probe, so the leader would correctly decline to reap while every joiner reaped on the leader's own observation; a mutation reproduces exactly that. The docs sentence promising LOCALAI_WORKER_TUNNEL=false restores direct dialling is gone. There is no such path, so it said the operator could take a worker dark and call it a rollback. Replaced with the upgrade order that is actually safe. The deadline spec the reviewer found vacuous now waits on the dial context's own Done channel before touching the stream, so the armed deadline has really expired; the mutation that survived for the reviewer reddens it. Nine mutations, each reddening a named spec, including both halves of isAbsenceClaim independently. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/distributed.go | 10 +- core/http/endpoints/localai/nodes.go | 35 ++-- core/services/cluster/dialer.go | 192 ++++++++++++++---- core/services/cluster/dialer_test.go | 129 +++++++++--- core/services/cluster/relay.go | 16 ++ .../nodes/backend_client_factory_test.go | 32 +++ core/services/nodes/file_stager_http.go | 9 + core/services/nodes/health.go | 14 ++ core/services/nodes/health_mock_test.go | 10 + core/services/nodes/health_test.go | 49 +++++ core/services/nodes/interfaces.go | 100 ++++++++- core/services/nodes/probe_cache.go | 29 ++- core/services/nodes/reconciler.go | 19 +- core/services/nodes/reconciler_prober_test.go | 108 ++++++++++ core/services/nodes/router.go | 31 ++- .../nodes/router_unreachable_worker_test.go | 120 ++++++++++- docs/content/features/distributed-mode.md | 13 +- pkg/grpc/backend.go | 21 +- pkg/grpc/client.go | 45 ++++ pkg/model/loader.go | 32 +++ pkg/model/remote_unroutable_internal_test.go | 61 ++++++ 21 files changed, 969 insertions(+), 106 deletions(-) create mode 100644 core/services/nodes/reconciler_prober_test.go create mode 100644 pkg/model/remote_unroutable_internal_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 512b81780296..fbe46a981212 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -375,10 +375,12 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade 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 + // An empty HTTPAddress is no longer a refusal. A tunnel-only worker + // reports none and does not need one: the http stream tag ignores + // the target and the worker routes to its own server. The host is + // only ever the URL's host component here, and WorkerHTTPHost + // supplies one that resolves nowhere so it cannot become a dial. + return nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), nil }, cfg.Distributed.RegistrationToken, workerHTTPDialer) xlog.Info("File stager initialized (HTTP direct transfer)") } diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index f10cc9c83ca9..f404555ac064 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -801,11 +801,11 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { - return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) - } - - resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, node.HTTPAddress, "/v1/backend-logs", registrationToken) + // No HTTPAddress guard: a tunnel-only worker reports none, and the + // http stream tag ignores the target anyway. WorkerHTTPHost fills the + // URL's host with something that identifies the node and resolves + // nowhere; the tunnel decides where the bytes go. + resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), "/v1/backend-logs", registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -831,12 +831,8 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { - 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, nodeID, node.HTTPAddress, path, registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), path, registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -888,7 +884,7 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s // NetDialContext below is what decides where the connection goes. A // missing dialer is a failure, not a direct dial: see // nodes.ErrNoWorkerDialer. - 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(nodeID, node.HTTPAddress), url.PathEscape(modelID)) workerHeaders := http.Header{} if registrationToken != "" { workerHeaders.Set("Authorization", "Bearer "+registrationToken) @@ -899,8 +895,21 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s workerDial = dialFor(nodeID) } if workerDial == nil { - return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, - fmt.Sprintf("cannot reach node %s: %v", nodeID, nodes.ErrNoWorkerDialer))) + // A JSON body cannot be written here: the response writer was + // hijacked by the upgrade above, so the status line is long gone + // and the write lands nowhere. The browser has to be told the same + // way every other failure past the upgrade tells it, with a close + // frame, and the socket has to be closed or it leaks for the life + // of the process. + // Best-effort: the browser may already have gone, and there is + // nothing left to report the failure to either way. The CLOSE is + // what matters and it is unconditional. + _ = browserWS.WriteMessage(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "no route to worker")) + _ = browserWS.Close() + xlog.Error("Cannot stream backend logs: no way to reach the worker", + "node", nodeID, "error", nodes.ErrNoWorkerDialer) + return nil } workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second, NetDialContext: workerDial} diff --git a/core/services/cluster/dialer.go b/core/services/cluster/dialer.go index b4a04551a973..fbb95d583acc 100644 --- a/core/services/cluster/dialer.go +++ b/core/services/cluster/dialer.go @@ -34,6 +34,102 @@ type PeerOpener interface { // and core/services/nodes reclaims the models of a worker it believes absent. var ErrNoRelayPath = errors.New("cluster: this replica cannot relay to the owner of that worker") +// ErrNoRoute reports that this replica could not get a request to a worker's +// backend, and is the FIFTH condition this phase keeps apart. +// +// It says nothing about whether the worker exists or is running. A worker's +// PRESENCE is its heartbeat, which lives in core/services/nodes and which this +// package cannot see; what this package can see is whether a route exists right +// now, and those are different questions with different answers. A worker that +// is registered, heartbeating and serving models can be unroutable from here +// for a whole list of ordinary reasons: it has not dialled its tunnel yet after +// a frontend-first upgrade, the replica holding its tunnel is restarting, the +// ownership row is a moment stale, this replica has no peer mesh. +// +// Every failure to RESOLVE OR OPEN a route carries it, so a consumer that must +// not act on absence has exactly one check to make. The specific condition +// stays in the unwrap chain underneath for anyone that can act on it, with one +// deliberate exception: see routeFailure. +// +// It is not carried by a REFUSAL from the worker itself. A worker that answers +// is present by demonstration, and folding its answer into "no route" would +// throw away the one thing on this path that is real evidence. +var ErrNoRoute = errors.New("cluster: no route from this replica to that worker") + +// noRouteError reports a worker this replica cannot route to, keeping the cause +// in its message and OUT of its unwrap chain. +// +// Withholding the cause is the entire point, and it is the same guarantee +// unreachableError makes for peers. The causes this is built over are absence +// claims: ErrNoConnection ("no live replica holds this worker's tunnel") and +// ErrInstanceNotFound ("no such frontend replica"). Both are true statements +// about the CLUSTER and neither is a statement about the worker, but a consumer +// matching on them would read them as one, and the consequence is the +// catastrophe this phase is built around: a scheduler concludes a worker that +// is heartbeating and serving has gone away, and reclaims its models. +// +// The guarantee therefore belongs to the type. There is no path by which an +// absence sentinel gets out, so no call site can leak one. +type noRouteError struct { + nodeID string + cause error +} + +func (e *noRouteError) Error() string { + return fmt.Sprintf("cluster: no route from this replica to node %q: %v", e.nodeID, e.cause) +} + +// Unwrap reports only ErrNoRoute. The cause reaches a human through Error() and +// reaches no error-matching caller at all. +func (e *noRouteError) Unwrap() error { return ErrNoRoute } + +// routeFailure is the ONE place a Dial failure is turned into an error, and the +// one place the absence rule is expressed. +// +// The rule: an absence claim never reaches a caller, and everything else stays +// matchable. It is a single predicate in a single function on purpose. An +// earlier shape in this phase encoded one policy in two predicates, and +// reverting either left the suite green because the error reached the same +// answer down the other path; a rule whose correctness argument IS its mutation +// evidence cannot afford to be un-mutatable in pieces. Falsifying either half +// of isAbsenceClaim now reddens a named spec. +func routeFailure(nodeID string, cause error) error { + if isAbsenceClaim(cause) { + return &noRouteError{nodeID: nodeID, cause: cause} + } + return fmt.Errorf("reaching node %q: %w: %w", nodeID, ErrNoRoute, cause) +} + +// isAbsenceClaim reports whether an error asserts that something does not +// exist. Those are the errors routeFailure keeps out of the chain. +// +// Both are about the CLUSTER rather than about the worker. ErrNoConnection says +// no live replica holds the worker's tunnel; ErrInstanceNotFound says a peer +// replica is not in the deployment. Neither can be answered by this package +// with "and therefore the worker is gone", because this package does not know +// what a worker is beyond an id in a connection row. +func isAbsenceClaim(err error) bool { + return errors.Is(err, ErrNoConnection) || errors.Is(err, ErrInstanceNotFound) +} + +// isWorkerAnswer reports whether an error is the WORKER's own refusal, read off +// the reply it sent. +// +// Those three sentinels are the only ones ReadStreamReply produces from a frame +// the worker actually wrote. Everything else it returns is a failure to read +// one, which is the tunnel breaking rather than the worker speaking. +// +// A reply carrying a code this frontend does not recognise is deliberately NOT +// counted here, even though a worker did send it. Classifying it as an answer +// would let a newer worker's vocabulary be read by an older frontend as +// evidence about a backend, and the consequence of guessing wrong in that +// direction is a reaped replica; guessing wrong the other way costs a retry. +func isWorkerAnswer(err error) bool { + return errors.Is(err, ErrStreamTagUnknown) || + errors.Is(err, ErrStreamTargetUnavailable) || + errors.Is(err, ErrStreamRequestInvalid) +} + // dialHandshakeTimeout bounds the request/reply exchange that opens every // stream, when the caller stated no deadline of its own. // @@ -79,13 +175,18 @@ func NewWorkerDialer(tunnels *TunnelRegistry, peers PeerOpener) *WorkerDialer { // inference that is quiet for minutes, and a deadline left over from the // handshake would abort it. // -// The errors are kept apart on purpose and a caller may act on them -// differently. ErrNoConnection means no live replica holds this worker's -// tunnel, which is the one answer that means the worker is absent. ErrNotOwner -// means the routing was stale and re-resolving may find it. ErrPeerUnreachable -// means a replica would not answer, ErrNoRelayPath that none could be dialled, -// and the tunnelproto sentinels that the worker itself refused. Nothing here -// ever converts one of the others into absence. +// EVERY failure to resolve or open a route carries ErrNoRoute, and NO failure +// carries an absence sentinel. That pair is the contract, and it is what makes +// this safe to consume from a package that reclaims a worker's models when it +// decides the worker has gone: there is one check to make, and there is nothing +// to mistake for absence even if the caller makes none. +// +// Underneath the umbrella the conditions stay apart and a caller may act on +// them differently. ErrNotOwner means the routing was stale and re-resolving +// may find it; ErrPeerUnreachable means a replica would not answer; +// ErrNoRelayPath means none could be dialled. A refusal from the WORKER carries +// its own tunnelproto sentinel and no umbrella at all, because a worker that +// answers has demonstrated it is there. func (d *WorkerDialer) Dial(ctx context.Context, nodeID, tag, target string) (net.Conn, error) { stream, err := d.tunnels.Open(ctx, nodeID) if err == nil { @@ -93,10 +194,9 @@ func (d *WorkerDialer) Dial(ctx context.Context, nodeID, tag, target string) (ne } if !errors.Is(err, ErrNotOwner) { // The tunnel is held HERE and its session would not carry a stream. - // Reported as itself: answering ErrNotOwner would send the caller to - // resolve an owner that is this same replica, and answering absence - // would tell a scheduler to reclaim a worker that is attached. - return nil, err + // ErrNotOwner stays out of it: that answer would send the caller to + // resolve an owner which is this same replica. + return nil, routeFailure(nodeID, err) } return d.relay(ctx, nodeID, tag, target) } @@ -137,10 +237,12 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string) (n // come back as ErrNoConnection here. owner, _, err := d.tunnels.reg.Owner(ctx, nodeID) if err != nil { - // ErrNoConnection and database failures both pass through as - // themselves. This is the ONLY path by which this function can produce - // an absence error, and it produces it only when Owner did. - return nil, err + // ErrNoConnection is the ordinary answer here, and it is precisely the + // one that must not get out: it means no live replica holds this + // worker's tunnel, which a worker that has not dialled in yet produces + // on every single request while it sits there heartbeating and serving. + // routeFailure keeps it in the message and out of the chain. + return nil, routeFailure(nodeID, err) } if owner == d.tunnels.selfID { // The table names this replica and the registry above said the tunnel @@ -149,19 +251,20 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string) (n // would resolve the same owner and relay again. Reported as the routing // fact so the caller re-resolves, which terminates: the row is either // re-claimed by whoever holds the worker now, or swept. - return nil, fmt.Errorf("opening a stream to node %q: the connection row names this replica, which no longer holds the tunnel: %w", nodeID, ErrNotOwner) + 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, fmt.Errorf("opening a stream to node %q held by replica %q: %w", nodeID, owner, ErrNoRelayPath) + 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 { - // Whatever the pool said, unchanged in its unwrap chain: - // ErrPeerUnreachable, ErrInstanceNotFound for an owner swept since the - // lookup above, or ErrPoolClosed while this process shuts down. None of - // them is a statement about the WORKER, and none is converted into one. - return nil, fmt.Errorf("opening a stream to node %q through replica %q: %w", nodeID, owner, err) + // ErrPeerUnreachable and ErrPoolClosed keep their identity; the one + // case the pool can also produce, ErrInstanceNotFound for an owner + // swept between the lookup above and this dial, is an absence claim + // about the REPLICA and routeFailure withholds it. Either way nothing + // here is a statement about the worker. + return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) } // The caller's remaining time, stated so the owning replica can bound its @@ -169,14 +272,15 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string) (n // the owner falls back to without it. if err := WriteRelayRequest(stream, nodeID, remainingBudget(ctx)); err != nil { _ = stream.Close() - return nil, fmt.Errorf("naming node %q on a stream to replica %q: %w", nodeID, owner, err) + return nil, routeFailure(nodeID, fmt.Errorf("naming the node on a stream to replica %q: %w", owner, err)) } if err := ReadRelayReply(stream); err != nil { - // ReadRelayReply already separates a refusal (ErrNotOwner, - // ErrRelayUnavailable, ErrRelayRequestInvalid) from a failure to read - // one, and neither kind is ever an absence error. + // A refusal from the OWNING REPLICA, not from the worker. It says the + // owner would not relay, which is a route that does not exist, so the + // umbrella is right for all of them; ErrNotOwner, ErrRelayUnavailable + // and ErrRelayRequestInvalid stay in the chain underneath. _ = stream.Close() - return nil, fmt.Errorf("relaying to node %q through replica %q: %w", nodeID, owner, err) + return nil, routeFailure(nodeID, fmt.Errorf("through replica %q: %w", owner, err)) } return d.handshake(ctx, stream, nodeID, tag, target) } @@ -189,20 +293,27 @@ func (d *WorkerDialer) relay(ctx context.Context, nodeID, tag, target string) (n // session, and a frontend that retries would exhaust the worker's stream // budget rather than the worker's patience. func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, tag, target string) (net.Conn, error) { - if deadline, ok := handshakeDeadline(ctx); ok { - if err := stream.SetDeadline(deadline); err != nil { - _ = stream.Close() - return nil, fmt.Errorf("arming the handshake deadline for node %q: %w", nodeID, err) - } + 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 stream would not carry the request, so the tunnel broke under it. + // Nothing was asked of the worker and nothing was learned about it. _ = stream.Close() - return nil, fmt.Errorf("asking node %q for %q on %q: %w", nodeID, tag, target, err) + return nil, routeFailure(nodeID, fmt.Errorf("asking for %q on %q: %w", tag, target, err)) } if err := ReadStreamReply(stream); err != nil { _ = stream.Close() - return nil, fmt.Errorf("opening %q on node %q: %w", tag, nodeID, err) + if isWorkerAnswer(err) { + // The worker wrote a refusal, so it is connected and answering. + // This is the ONE failure on the whole path that is real evidence + // about the worker, and putting the umbrella on it would throw that + // 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 unconditionally rather than only when one was armed, so that this @@ -213,7 +324,7 @@ func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, t // bounds a peer that has stopped answering. if err := stream.SetDeadline(time.Time{}); err != nil { _ = stream.Close() - return nil, fmt.Errorf("clearing the handshake deadline for node %q: %w", nodeID, err) + 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 @@ -221,13 +332,18 @@ func (d *WorkerDialer) handshake(ctx context.Context, stream net.Conn, nodeID, t // handshakeDeadline is when the handshake must be done by: the caller's own // deadline when it has one and it is the sooner, and the backstop otherwise. -func handshakeDeadline(ctx context.Context) (time.Time, bool) { +// +// There is always one, which is why this returns no "was there one" flag: a +// context with no deadline still gets the backstop, so the caller has nothing +// to branch on. It used to return a bool that was unconditionally true, and the +// branch behind it could not be taken. +func handshakeDeadline(ctx context.Context) time.Time { backstop := time.Now().Add(dialHandshakeTimeout) deadline, ok := ctx.Deadline() if !ok || deadline.After(backstop) { - return backstop, true + return backstop } - return deadline, true + return deadline } // remainingBudget is how long the caller is still willing to wait, or zero when diff --git a/core/services/cluster/dialer_test.go b/core/services/cluster/dialer_test.go index ffcfe3775e13..2e0b945d3704 100644 --- a/core/services/cluster/dialer_test.go +++ b/core/services/cluster/dialer_test.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "net" + "sync" "time" "github.com/mudler/LocalAI/core/services/cluster" @@ -26,15 +27,31 @@ import ( type stubPeers struct { sess *yamux.Session err error + + // opened records the peers this pool was asked for, so a spec can assert + // that a replica was NOT dialled. That is the only way to tell a dialer + // that resolved a live owner from one that resolved a dead row and then + // found out the hard way. + mu sync.Mutex + opened []string } -func (s *stubPeers) Open(ctx context.Context, _ string) (net.Conn, error) { +func (s *stubPeers) Open(ctx context.Context, peerID string) (net.Conn, error) { + s.mu.Lock() + s.opened = append(s.opened, peerID) + s.mu.Unlock() if s.err != nil { return nil, s.err } return s.sess.OpenStream(ctx) } +func (s *stubPeers) peersDialled() []string { + s.mu.Lock() + defer s.mu.Unlock() + return append([]string(nil), s.opened...) +} + // dialResult carries what a Dial produced, so a spec can wait on a channel // rather than on a clock. type dialResult struct { @@ -124,7 +141,8 @@ func refuseOneStream(worker *yamux.Session, reason error) { // // It is the assertion this whole phase turns on. core/services/nodes reclaims a // worker's models when it concludes the worker is absent, so an unreachable -// peer or a refusing worker arriving as absence would evict healthy work. +// peer, a stale ownership row or a worker that has not dialled its tunnel yet +// arriving as absence would evict healthy work. func expectNotAbsence(err error) { GinkgoHelper() Expect(err).To(HaveOccurred()) @@ -132,6 +150,18 @@ func expectNotAbsence(err error) { Expect(err).ToNot(MatchError(cluster.ErrInstanceNotFound)) } +// expectNoRoute asserts the umbrella is present as well as absence being gone. +// +// The umbrella is what crosses the package boundary. A consumer that reclaims +// models has one check to make, and it can only make it if EVERY failure to +// resolve or open a route carries it; a single path that forgets is a path +// where a live worker gets reaped. +func expectNoRoute(err error) { + GinkgoHelper() + expectNotAbsence(err) + Expect(err).To(MatchError(cluster.ErrNoRoute)) +} + var _ = Describe("The worker dialer", func() { var ( db *gorm.DB @@ -200,16 +230,25 @@ var _ = Describe("The worker dialer", func() { Expect(string(echoed)).To(Equal("ping")) }) - It("leaves no read deadline armed on the stream it hands back", func() { + It("leaves no deadline armed on the stream it hands back", func() { // The handshake is bounded; the request that follows it is the // caller's business and may be a generation that is quiet for - // minutes. A deadline left armed here would abort it. + // minutes. A deadline left armed here would abort it, and in + // production the dial context is the model-load or request budget, + // so the stream would die tens of seconds in. + // + // The first version of this spec did not assert that. It set a + // 300ms context and then wrote immediately, so the armed deadline + // had not expired and deleting the clear left it green: it detected + // only a deadline set in the PAST. What makes it bite is waiting for + // the dial context to actually expire FIRST, on its own Done channel + // rather than a sleep, and only then using the stream. frontend, worker := workerTunnel() _, err := mine.Attach(ctx, "w1", frontend) Expect(err).ToNot(HaveOccurred()) seen := serveOneStream(worker) - deadlined, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + deadlined, cancel := context.WithTimeout(ctx, 200*time.Millisecond) defer cancel() d := cluster.NewWorkerDialer(mine, nil) result := dialAsync(d, deadlined, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000") @@ -220,12 +259,16 @@ var _ = Describe("The worker dialer", func() { Expect(out.err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = out.conn.Close() }) - // The dial context's deadline has now passed. A stream still - // carrying it would fail this write and this read. - Eventually(func() error { - _, err := out.conn.Write([]byte("ping")) - return err - }, "10s").Should(Succeed()) + // The one wait this spec cannot replace with an event of its own: + // there is nothing to observe until the dial's deadline is behind + // us, and the deadline is the thing under test. + <-deadlined.Done() + Expect(deadlined.Err()).To(HaveOccurred()) + + // Both directions, because SetDeadline arms read and write and a + // clear that only covered one would still kill a live request. + _, err = out.conn.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) echoed := make([]byte, 4) Eventually(readInto(out.conn, echoed), "10s").Should(Receive(BeNil())) Expect(string(echoed)).To(Equal("ping")) @@ -241,8 +284,12 @@ var _ = Describe("The worker dialer", func() { var out dialResult Eventually(dialAsync(d, ctx, "w1", "nonsense", ""), "10s").Should(Receive(&out)) Expect(out.err).To(MatchError(cluster.ErrStreamTagUnknown)) - // A refusal is PROOF the worker is connected and answered. + // A refusal is PROOF the worker is connected and answered, so it is + // the ONE failure on this path that carries no umbrella: it is real + // evidence about the worker, and folding it into "no route" would + // throw that evidence away. expectNotAbsence(out.err) + Expect(out.err).ToNot(MatchError(cluster.ErrNoRoute)) }) It("reports a broken tunnel held here as itself, not as a routing fact", func() { @@ -257,9 +304,8 @@ var _ = Describe("The worker dialer", func() { d := cluster.NewWorkerDialer(mine, nil) var out dialResult Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) - Expect(out.err).To(HaveOccurred()) Expect(out.err).ToNot(MatchError(cluster.ErrNotOwner)) - expectNotAbsence(out.err) + expectNoRoute(out.err) }) }) @@ -376,7 +422,7 @@ var _ = Describe("The worker dialer", func() { var out dialResult Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "20s").Should(Receive(&out)) Expect(out.err).To(MatchError(cluster.ErrPeerUnreachable)) - expectNotAbsence(out.err) + expectNoRoute(out.err) }) It("passes a stale ownership refusal back as the routing fact", func() { @@ -396,7 +442,7 @@ var _ = Describe("The worker dialer", func() { var out dialResult Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) Expect(out.err).To(MatchError(cluster.ErrNotOwner)) - expectNotAbsence(out.err) + expectNoRoute(out.err) }) It("refuses rather than relaying to itself when the table names this replica", func() { @@ -410,7 +456,7 @@ var _ = Describe("The worker dialer", func() { var out dialResult Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) Expect(out.err).To(MatchError(cluster.ErrNotOwner)) - expectNotAbsence(out.err) + expectNoRoute(out.err) }) It("reports having no way to relay as its own condition", func() { @@ -424,22 +470,35 @@ var _ = Describe("The worker dialer", func() { Expect(out.err).To(MatchError(cluster.ErrNoRelayPath)) Expect(out.err).ToNot(MatchError(cluster.ErrNotOwner)) Expect(out.err).ToNot(MatchError(cluster.ErrPeerUnreachable)) - expectNotAbsence(out.err) + expectNoRoute(out.err) }) }) Describe("when no live replica holds the tunnel", func() { - It("reports absence when the worker has no connection row at all", func() { - d := cluster.NewWorkerDialer(mine, &stubPeers{err: errors.New("no peer should be dialled")}) + // The rolling-upgrade case, and the one this phase must not get wrong. + // + // A worker's PRESENCE is its heartbeat, which lives in + // core/services/nodes. "No live replica holds this worker's tunnel" is + // a fact about tunnels and says nothing about the worker: a worker that + // has not dialled in yet after a frontend-first upgrade produces it on + // every request while it sits there heartbeating and serving models. + // A consumer told that is absence reclaims every one of those models. + It("answers no-route, never absence, for a worker with no connection row", func() { + peers := &stubPeers{err: errors.New("no peer should be dialled")} + d := cluster.NewWorkerDialer(mine, peers) var out dialResult Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) - Expect(out.err).To(MatchError(cluster.ErrNoConnection)) + expectNoRoute(out.err) + // The cause still reaches a human. + Expect(fmt.Sprint(out.err)).To(ContainSubstring("no connection recorded")) + Expect(peers.peersDialled()).To(BeEmpty()) }) - It("reports absence when the row's owner has stopped heartbeating", func() { + It("answers no-route, never absence, when the row's owner has stopped heartbeating", func() { // End to end over the join Owner does: the row is there, the owner - // is not. A dialer built on the unjoined read would dial a corpse - // and report the worker as unreachable rather than as absent. + // is not. The join is what stops this replica dialling a process + // that is gone, which is why the spec asserts no peer was dialled + // as well as what came back. Expect(reg.Register(ctx, "ghost", "10.0.0.9:8080", "v1")).To(Succeed()) _, err := reg.Claim(ctx, "w1", "ghost") Expect(err).ToNot(HaveOccurred()) @@ -451,10 +510,28 @@ var _ = Describe("The worker dialer", func() { Expect(err).ToNot(HaveOccurred()) Expect(owner).To(Equal("ghost")) - d := cluster.NewWorkerDialer(mine, &stubPeers{err: errors.New("no peer should be dialled")}) + peers := &stubPeers{err: errors.New("no peer should be dialled")} + d := cluster.NewWorkerDialer(mine, peers) + var out dialResult + Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) + expectNoRoute(out.err) + Expect(peers.peersDialled()).To(BeEmpty()) + }) + + It("keeps an owner swept mid-dial out of the chain as well", func() { + // The other absence sentinel. PeerPool resolves the owner's address + // through the registry, so a replica reaped between Owner and the + // dial comes back as ErrInstanceNotFound. That is absence of a + // REPLICA, and a consumer matching absence would read it as absence + // of the WORKER. + Expect(reg.Register(ctx, "owner", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "owner") + Expect(err).ToNot(HaveOccurred()) + + d := cluster.NewWorkerDialer(mine, &stubPeers{err: cluster.ErrInstanceNotFound}) var out dialResult Eventually(dialAsync(d, ctx, "w1", cluster.StreamTagGRPC, "127.0.0.1:41000"), "10s").Should(Receive(&out)) - Expect(out.err).To(MatchError(cluster.ErrNoConnection)) + expectNoRoute(out.err) }) }) diff --git a/core/services/cluster/relay.go b/core/services/cluster/relay.go index 9392834f4e14..61594c7abecb 100644 --- a/core/services/cluster/relay.go +++ b/core/services/cluster/relay.go @@ -125,6 +125,16 @@ func ReadRelayRequest(r io.Reader) (string, time.Duration, error) { if err != nil { return "", 0, fmt.Errorf("reading a relay request for node %q: budget %q is not a number of milliseconds: %w", nodeID, budgetText, err) } + if millis > maxRelayBudgetMillis { + // time.Duration is nanoseconds in an int64, so multiplying by + // time.Millisecond overflows past about 2.9e11 ms. Overflow here is + // bounded in the safe direction (it can only produce a negative or a + // small value, and both shorten the declaring peer's OWN open), but a + // bound that holds by arithmetic accident is not a bound. Anything past + // the ceiling is clamped to it, because a caller claiming to wait + // longer than the relay's own backstop gets the backstop either way. + millis = maxRelayBudgetMillis + } if millis <= 0 { // A caller with nothing left to spend. Reported as such rather than // folded into "not stated", so the relay refuses at once instead of @@ -209,6 +219,12 @@ func ReadRelayReply(r io.Reader) error { } } +// maxRelayBudgetMillis is the largest budget a peer may declare, and exists so +// the conversion below cannot overflow. A day is many orders of magnitude past +// relayOpenTimeout, which is the only thing a budget is ever compared against, +// so clamping to it changes no honest caller's behaviour. +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 killed between OpenStream and diff --git a/core/services/nodes/backend_client_factory_test.go b/core/services/nodes/backend_client_factory_test.go index be20f5e347b0..15adeb4f3117 100644 --- a/core/services/nodes/backend_client_factory_test.go +++ b/core/services/nodes/backend_client_factory_test.go @@ -5,6 +5,7 @@ package nodes import ( "context" "net" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -99,3 +100,34 @@ var _ = Describe("The backend client factory", func() { }) }) }) + +var _ = Describe("the host used to address a worker's own HTTP server", func() { + It("uses the registered address when the worker reports one", func() { + Expect(WorkerHTTPHost("node-1", "10.0.0.5:8080")).To(Equal("10.0.0.5:8080")) + }) + + It("still produces a host for a tunnel-only worker that reports none", func() { + // Task 7 removes the worker's inbound listeners, at which point a + // worker has no address to report. Refusing here would refuse exactly + // the workers the tunnel exists for, and the guards that used to do + // that returned 502 "node has no HTTP address". + host := WorkerHTTPHost("node-1", "") + Expect(host).ToNot(BeEmpty()) + Expect(host).To(ContainSubstring("node-1")) + }) + + It("produces a host that cannot resolve, so it can never become a dial", func() { + // The value fills a URL's host component and nothing else. Making it + // unresolvable is what stops a later refactor connecting to it by + // accident: .invalid is reserved by RFC 2606 and resolves nowhere. + host := WorkerHTTPHost("node-1", "") + hostname, _, err := net.SplitHostPort(host) + Expect(err).ToNot(HaveOccurred()) + Expect(hostname).To(HaveSuffix(".invalid")) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _, err = net.DefaultResolver.LookupHost(ctx, hostname) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/nodes/file_stager_http.go b/core/services/nodes/file_stager_http.go index bd286e5d3426..825560c82df4 100644 --- a/core/services/nodes/file_stager_http.go +++ b/core/services/nodes/file_stager_http.go @@ -41,6 +41,15 @@ type HTTPFileStager struct { // clients caches one *http.Client per node. Caching is what keeps the // connection pool: a client built per request would open a fresh tunnel // stream for every chunk of a multi-gigabyte upload. + // + // Entries are never pruned, and that is judged acceptable rather than + // overlooked. The map is bounded by the number of distinct workers this + // frontend has ever staged to, which is bounded by the fleet; each entry is + // a transport whose idle connections the 90s IdleConnTimeout above reclaims, + // so a departed worker's entry holds a map slot and nothing else. It is the + // same shape as PeerPool.links and would need the same thing to fix + // properly: a signal that a node has left, which the deregistration path + // does not publish today. clientsMu sync.Mutex clients map[string]*http.Client responseTimeout time.Duration // timeout waiting for server response after upload diff --git a/core/services/nodes/health.go b/core/services/nodes/health.go index 44d79e0d8ab3..8a26a042f48c 100644 --- a/core/services/nodes/health.go +++ b/core/services/nodes/health.go @@ -203,9 +203,23 @@ func (hm *HealthMonitor) doCheckAll(ctx context.Context) { mCheckCtx, mCancel := context.WithTimeout(ctx, 5*time.Second) ok, _ := mClient.HealthCheck(mCheckCtx) mCancel() + // Asked BEFORE the client is closed, because closing is what + // would discard the transport's record of why it failed. + unreached := unroutable(mClient) if closer, ok := mClient.(io.Closer); ok { closer.Close() } + if unreached != nil { + // The probe never reached a backend, so it observed + // nothing. The miss streak is left exactly as it was: + // neither advanced, which after three passes would delete + // this row and every other row in the fleet the moment a + // peer link blipped, nor cleared, which would forgive a + // backend that really has died. + xlog.Warn("Could not probe a model backend: no route to the worker", + "node", node.ID, "model", m.ModelName, "replica", m.ReplicaIndex, "error", unreached) + 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 281b876a7e4c..737be51b1765 100644 --- a/core/services/nodes/health_mock_test.go +++ b/core/services/nodes/health_mock_test.go @@ -133,8 +133,17 @@ func (f *fakeNodeHealthStore) RemoveNodeModel(_ context.Context, nodeID, modelNa type fakeBackendClient struct { healthy bool err error + // dialErr makes this client report that its TRANSPORT failed, which is what + // a real client whose tunnel dial failed does. It is the half of + // unroutability that a refusing factory cannot stand in for, and the likely + // one in production: the factory only fails when the wiring is absent. + dialErr error } +// LastDialError satisfies grpc.DialErrorReporter so a spec can drive the +// "reached no backend" branch without a real tunnel. +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 @@ -391,4 +400,5 @@ func freshTime() time.Time { // Compile-time interface checks var _ NodeHealthStore = (*fakeNodeHealthStore)(nil) var _ BackendClientFactory = (*fakeBackendClientFactory)(nil) +var _ grpc.DialErrorReporter = (*fakeBackendClient)(nil) var _ grpc.Backend = (*fakeBackendClient)(nil) diff --git a/core/services/nodes/health_test.go b/core/services/nodes/health_test.go index 6f624b0cacb5..a20384737f94 100644 --- a/core/services/nodes/health_test.go +++ b/core/services/nodes/health_test.go @@ -9,6 +9,8 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" "gorm.io/gorm" ) @@ -325,6 +327,53 @@ var _ = Describe("HealthMonitor (mock-based)", func() { Expect(store.getNode("node-cut").Status).To(Equal(StatusHealthy)) }) + It("leaves a model row alone when the probe never reached the worker", func() { + // The sibling of the factory case above, and the likelier one. The + // client is built fine and the tunnel DIAL fails, which gRPC + // reports with the same code as a dead backend. Counted as a miss + // it would delete every model row in the fleet after three passes + // of a peer link blip, while the models kept serving. + store := newFakeNodeHealthStore() + factory := newFakeBackendClientFactory() + hm := newTestHealthMonitor(store, factory, true, staleThreshold) + hm.perModelHealthCheck = true + + node := makeTestNode("node-blip", "blip-worker", "10.0.0.22:50051", StatusHealthy, freshTime()) + store.addNode(node) + store.addNodeModel("node-blip", NodeModel{NodeID: "node-blip", ModelName: "m", Address: "10.0.0.22:50053"}) + factory.setClient("10.0.0.22:50053", &fakeBackendClient{ + healthy: false, + err: fmt.Errorf("connection error"), + dialErr: fmt.Errorf("%w: %w", cluster.ErrNoRoute, cluster.ErrPeerUnreachable), + }) + + for i := 0; i < perModelMissThreshold+2; i++ { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).NotTo(ContainElement(ContainSubstring("RemoveNodeModel"))) + }) + + It("still reaps a backend that died on a worker it CAN reach", func() { + // The other direction, so the new check cannot pass by never + // reaping. A dial that succeeded and an RPC that failed is a dead + // process, and its row must still go. + store := newFakeNodeHealthStore() + factory := newFakeBackendClientFactory() + hm := newTestHealthMonitor(store, factory, true, staleThreshold) + hm.perModelHealthCheck = true + + node := makeTestNode("node-dead", "dead-worker", "10.0.0.23:50051", StatusHealthy, freshTime()) + store.addNode(node) + store.addNodeModel("node-dead", NodeModel{NodeID: "node-dead", ModelName: "m", Address: "10.0.0.23:50053"}) + // No dialErr: the transport was fine. + factory.setClient("10.0.0.23:50053", &fakeBackendClient{healthy: false, err: fmt.Errorf("connection refused")}) + + for i := 0; i < perModelMissThreshold; i++ { + hm.doCheckAll(context.Background()) + } + Expect(store.getCalls()).To(ContainElement("RemoveNodeModel:node-dead:m:0")) + }) + It("preserves model row when an intermittent failure is followed by a success", func() { store := newFakeNodeHealthStore() factory := newFakeBackendClientFactory() diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 1231c811f371..d1028983843b 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -143,10 +143,15 @@ type NodeManager interface { // WorkerDialerFor hands back the dial function for one worker's backend // processes: the shape grpc.WithContextDialer wants, bound to a node. // -// It is declared here rather than taken as a core/services/cluster type so that -// this package keeps no dependency on that one. cluster is a leaf and imports -// neither this package nor core/http; the wiring in core/application supplies -// (*cluster.WorkerDialer).GRPCDialerFor, which has exactly this shape. +// A function type rather than a *cluster.WorkerDialer, so nothing here is bound +// to that concrete type and a spec can supply a dial without building a tunnel +// registry, a peer pool and a database. It is NOT to avoid a dependency: this +// package already imports core/services/cluster (registry.go, for Migrate), and +// an earlier version of this comment claimed otherwise. The dependency that +// does matter runs the other way, and cluster is held to it by go list -deps. +// +// core/application supplies (*cluster.WorkerDialer).GRPCDialerFor, which has +// exactly this shape. type WorkerDialerFor func(nodeID string) func(ctx context.Context, addr string) (net.Conn, error) // WorkerNetDialerFor hands back the dial function for one worker's own HTTP @@ -155,6 +160,26 @@ type WorkerDialerFor func(nodeID string) func(ctx context.Context, addr string) // has this shape. type WorkerNetDialerFor func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) +// ErrWorkerUnroutable reports that this frontend could not get a request to a +// worker's backend, and says NOTHING about whether that worker or its backend +// is alive. +// +// It is the fifth condition, on this side of the package boundary. A worker's +// presence is its HEARTBEAT, and this package owns that; a route to it is a +// separate fact owned by core/services/cluster, and the two now differ. A +// worker can be registered, heartbeating and serving every request another +// replica sends it while being unroutable from here: it has not dialled its +// tunnel yet after a frontend-first upgrade, the replica holding its tunnel is +// restarting, the ownership row is a moment stale, this replica has no peer +// mesh. Every one of those used to be indistinguishable from "the backend +// process died", because gRPC reports both as codes.Unavailable. +// +// Everything in this package that DELETES a node_models row must consult it +// first. That is the phase's stated catastrophe in its concrete form: a row +// deleted here is a model reclaimed and reloaded elsewhere, so mistaking a peer +// link blip for a dead backend evicts healthy work across the fleet at once. +var ErrWorkerUnroutable = errors.New("nodes: this frontend has no route to that worker") + // ErrNoWorkerDialer reports that something tried to reach a worker without a // way to reach it through the worker's tunnel. // @@ -162,9 +187,41 @@ type WorkerNetDialerFor func(nodeID string) func(ctx context.Context, network, a // advertised address. A worker that holds a tunnel need not listen on anything // and may be behind NAT with no address to dial, so the fallback would work // only where the tunnel was not needed: on a single-host developer setup, and -// nowhere the feature exists for. It is also not an absence error: nothing -// about it says the worker is gone. -var ErrNoWorkerDialer = errors.New("nodes: no worker tunnel dialer is configured, so this worker cannot be reached") +// nowhere the feature exists for. +// +// It is a SPECIALISATION of ErrWorkerUnroutable rather than a sibling, so the +// single check every reaping path makes covers both. The difference between +// them is only when they happen: this one is a boot-time misconfiguration, and +// the general one is a running deployment losing a route for a moment. Neither +// is a statement about the worker. +var ErrNoWorkerDialer = fmt.Errorf("%w: no worker tunnel dialer is configured", ErrWorkerUnroutable) + +// unroutable reports why a call on client never reached the backend, or nil +// when it did reach one. +// +// This is where core/services/cluster's five conditions cross the package +// boundary. They cannot cross on the RPC error: gRPC turns any dialer failure +// into codes.Unavailable with the cause flattened into a message, and +// codes.Unavailable is ALSO what a backend process that has died produces. +// pkg/grpc records the dialer's error VALUE instead, so cluster.ErrNoRoute and +// whatever sits under it are still matchable here. +// +// A client that reports nothing (no custom dialer, or a test double) yields +// nil, which means "the call reached a backend" and preserves the behaviour +// every non-distributed caller has always had. +func unroutable(client grpc.Backend) error { + reporter, ok := client.(grpc.DialErrorReporter) + if !ok { + return nil + } + dialErr := reporter.LastDialError() + if dialErr == nil { + return nil + } + // Multi-%w: the umbrella this package acts on, and the cluster condition + // underneath it, both stay matchable. + return fmt.Errorf("%w: %w", ErrWorkerUnroutable, dialErr) +} // BackendClientFactory creates the gRPC clients this frontend uses to reach // model backends running on worker nodes. @@ -233,3 +290,32 @@ func (f *tunnelClientFactory) NewClientForNode(nodeID, address string, parallel } return grpc.NewClientWithDialer(address, parallel, nil, false, f.token, dial), nil } + +// unroutableHostSuffix is appended to a node id to build a Host for a worker +// that reports no HTTP address. +// +// .invalid is reserved by RFC 2606 and resolves nowhere, which is the point: +// the string exists ONLY to fill the host component of a URL, and a value that +// could resolve would be one a future refactor could accidentally connect to. +const unroutableHostSuffix = ".worker.invalid:80" + +// WorkerHTTPHost is the host to put in a URL addressed to a worker's own HTTP +// server. +// +// A tunnel-only worker has no inbound address to report, and after this phase +// it does not need one: the `http` stream tag ignores the target entirely and +// the worker routes the stream to its own server wherever that bound. But an +// http.Request still needs a host, so refusing an empty HTTPAddress would +// refuse exactly the workers the tunnel exists for. This returns a name that +// identifies the node for logs and for the Host header, and that nothing can +// connect to. +// +// It is NOT a dial target and never becomes one. Every caller pairs it with a +// transport whose DialContext is that node's tunnel, so the host is read and +// discarded; see cluster.WorkerDialer.DialerFor and the `http` tag. +func WorkerHTTPHost(nodeID, httpAddress string) string { + if httpAddress != "" { + return httpAddress + } + return nodeID + unroutableHostSuffix +} diff --git a/core/services/nodes/probe_cache.go b/core/services/nodes/probe_cache.go index 422e36ede4e2..fee71383e7e1 100644 --- a/core/services/nodes/probe_cache.go +++ b/core/services/nodes/probe_cache.go @@ -73,22 +73,41 @@ 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 { + alive, _ := c.DoOrCachedResult(key, func() (bool, error) { return probe(), nil }) + return alive +} + +// DoOrCachedResult is DoOrCached with a second result: the reason the probe +// never reached the backend, or nil when it did. +// +// The second result travels through the SINGLEFLIGHT, which is the whole reason +// it is not simply a variable the caller closes over. A closed-over variable is +// only written by the goroutine that actually runs the probe; every other +// caller coalesced into that flight adopts the leader's boolean and sees its own +// unset variable, so the leader would correctly decline to reap while its +// joiners reaped on the very same observation. Carrying it in singleflight's +// error slot hands every joiner the leader's reason as well as its answer. +// +// A probe that could not reach the backend is NOT cached either way. Caching it +// as fresh would hide a genuinely dead backend behind a network blip, and +// caching it as a failure is what Invalidate already does. +func (c *probeCache) DoOrCachedResult(key string, probe func() (bool, error)) (bool, error) { if c.IsFresh(key) { - return true + return true, nil } - v, _, _ := c.flight.Do(key, func() (any, error) { + v, unreached, _ := 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 } - ok := probe() + ok, unreached := probe() if ok { c.markFresh(key) } else { c.Invalidate(key) } - return ok, nil + return ok, unreached }) - return v.(bool) + return v.(bool), unreached } diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index d50b55cef716..923559932d16 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -91,17 +91,34 @@ 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 unreached := unroutable(client); unreached != nil { + // The RPC never reached a backend. classifyProbeOutcome cannot tell: + // gRPC hands it codes.Unavailable for a worker this frontend has no + // route to and for a backend process that has died, and the reaper + // deletes rows on the second. + xlog.Warn("Could not probe a model: no route to the worker", + "node", nodeID, "address", address, "error", unreached) + return ProbeUnknown + } return classifyProbeOutcome(ok, err) } // classifyProbeOutcome maps a HealthCheck result onto a ProbeOutcome. // +// It is only ever reached for a probe that DID reach the worker. That is a +// precondition and not an observation it can make for itself: its caller asks +// the transport first and answers ProbeUnknown when the dial failed. Without +// that step the Unavailable case below is wrong, because a worker this frontend +// cannot route to produces exactly the same code as a backend that has died, +// and only one of the two should cost a row. +// // The gRPC client is lazy, so connection failures surface on the RPC rather // than at dial time, and the status code tells the two cases apart: // // - DeadlineExceeded: the transport was fine but nothing serviced the RPC in // time. That is a backend stuck inside a long synchronous request. -// - Unavailable: nothing is listening. The process is gone. +// - Unavailable: the worker was reached and nothing is listening on that +// port. The process is gone. // // A blackholed network also yields DeadlineExceeded and is therefore treated as // busy. That is deliberate: whole-node failures are the health monitor's job diff --git a/core/services/nodes/reconciler_prober_test.go b/core/services/nodes/reconciler_prober_test.go new file mode 100644 index 000000000000..2bacef554297 --- /dev/null +++ b/core/services/nodes/reconciler_prober_test.go @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "errors" + "fmt" + "net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +// proberFactory hands the prober one client, and records what it was asked for. +type proberFactory struct { + client grpc.Backend + err error + asked []string +} + +func (f *proberFactory) NewClientForNode(nodeID, address string, _ bool) (grpc.Backend, error) { + f.asked = append(f.asked, nodeID+"|"+address) + if f.err != nil { + return nil, f.err + } + return f.client, nil +} + +var _ = Describe("the reconciler's gRPC model prober", func() { + // The two lines that decide whether a row survives, both previously + // untested. Everything else in the reaper is driven through fakeProber, + // which means the mapping from a real client to a ProbeOutcome had nothing + // holding it at all. + probe := func(f *proberFactory) ProbeOutcome { + GinkgoHelper() + return grpcModelProber{clients: f}.Probe(context.Background(), "node-1", "10.0.0.1:9001") + } + + It("answers ProbeUnknown when no client can be built for the node", func() { + Expect(probe(&proberFactory{err: ErrNoWorkerDialer})).To(Equal(ProbeUnknown)) + }) + + It("answers ProbeUnknown when the client was built and the tunnel dial failed", func() { + // The likelier half. ProbeUnreachable here would delete the row after + // probeFailuresBeforeReap passes of a peer link that was merely + // restarting, and the backend would still be running the model. + Expect(probe(&proberFactory{client: &fakeBackendClient{ + healthy: false, + err: fmt.Errorf("rpc error: code = Unavailable"), + dialErr: fmt.Errorf("%w: %w", cluster.ErrNoRoute, cluster.ErrPeerUnreachable), + }})).To(Equal(ProbeUnknown)) + }) + + It("asks for the client by NODE, not by address alone", func() { + f := &proberFactory{client: &fakeBackendClient{healthy: true}} + Expect(probe(f)).To(Equal(ProbeAlive)) + Expect(f.asked).To(ContainElement("node-1|10.0.0.1:9001")) + }) + + It("answers ProbeAlive for a healthy backend it reached", func() { + Expect(probe(&proberFactory{client: &fakeBackendClient{healthy: true}})).To(Equal(ProbeAlive)) + }) + + It("still answers ProbeUnreachable for a dead backend on a worker it reached", func() { + // The other direction. The new check must not turn the reaper off: a + // backend that answered "unhealthy" over a working transport is a ghost + // and its row should go. + Expect(probe(&proberFactory{client: &fakeBackendClient{healthy: false}})).To(Equal(ProbeUnreachable)) + }) + + It("does not report a transport that recovered", func() { + // LastDialError is cleared by a successful dial, so a client that + // failed once and then reconnected must not keep reading as + // unroutable; otherwise a row could never be reaped again after one + // blip on that client. + listener, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = listener.Close() }) + + attempt := 0 + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + var d net.Dialer + return func(ctx context.Context, _ string) (net.Conn, error) { + attempt++ + if attempt == 1 { + return nil, errors.New("first dial fails") + } + return d.DialContext(ctx, "tcp", listener.Addr().String()) + } + }) + Expect(err).ToNot(HaveOccurred()) + client, err := f.NewClientForNode("node-1", "10.0.0.1:9001", false) + Expect(err).ToNot(HaveOccurred()) + + _, _ = client.HealthCheck(context.Background()) + Expect(unroutable(client)).ToNot(BeNil()) + + // gRPC re-dials on the next call; the listener now accepts. + Eventually(func() error { + _, _ = client.HealthCheck(context.Background()) + return unroutable(client) + }, "20s").Should(BeNil()) + }) +}) diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index b9731f4a305a..2845d7f9898b 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -1941,11 +1941,18 @@ func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir // (DecrementInFlight + RemoveNodeModel) still triggers on the next request. // // The client is built OUTSIDE the memoized closure, which is what keeps an -// unreachable worker out of the cache entirely: DoOrCached only ever sees a -// real answer. Building it costs a struct and no I/O, since the gRPC client +// unreachable worker out of the cache entirely: DoOrCachedResult only ever sees +// a real answer. Building it costs a struct and no I/O, since the gRPC client // dials lazily on its first call. +// +// The client is the RAW factory client rather than buildClientForAddr's, on +// purpose. A health check stages no files, so the staging wrapper buys nothing +// here; and the wrapper hides the transport, because it embeds grpc.Backend and +// so does not carry LastDialError through. Wrapping would leave this function +// unable to tell a dead backend from an unreachable worker, which is the whole +// question it now answers. func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr string) (alive, probed bool) { - client, err := r.buildClientForAddr(node, addr, false) + client, err := r.clientFactory.NewClientForNode(node.ID, addr, false) if err != nil { xlog.Error("Cannot probe a model backend: no way to reach the worker", "node", node.ID, "address", addr, "error", err) @@ -1954,12 +1961,24 @@ func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr s defer closeClient(client) key := node.ID + "|" + addr - return r.probeCache.DoOrCached(key, func() bool { + alive, unreached := r.probeCache.DoOrCachedResult(key, func() (bool, error) { checkCtx, cancel := context.WithTimeout(ctx, 2*time.Second) defer cancel() ok, _ := client.HealthCheck(checkCtx) - return ok - }), true + if ok { + return true, nil + } + // The RPC failed. gRPC reports a dead backend and an unreachable + // worker with the same code, so the only way to tell them apart is to + // ask the transport whether it was the one that failed. + return false, unroutable(client) + }) + if unreached != nil { + xlog.Warn("Could not probe a model backend: no route to the worker", + "node", node.ID, "address", addr, "error", unreached) + return false, false + } + return alive, true } // closeClient closes a gRPC backend client if it implements io.Closer. diff --git a/core/services/nodes/router_unreachable_worker_test.go b/core/services/nodes/router_unreachable_worker_test.go index 943741396e11..206956e1194f 100644 --- a/core/services/nodes/router_unreachable_worker_test.go +++ b/core/services/nodes/router_unreachable_worker_test.go @@ -4,20 +4,38 @@ package nodes import ( "context" - "errors" + "fmt" + "net" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/services/cluster" grpc "github.com/mudler/LocalAI/pkg/grpc" ) // unreachableClientFactory cannot build a client for any node, standing in for -// a frontend whose worker tunnel dialer is missing or broken. +// a frontend whose worker tunnel dialer is missing or broken. This is the +// BOOT-TIME half of unroutability. type unreachableClientFactory struct{} func (unreachableClientFactory) NewClientForNode(_, _ string, _ bool) (grpc.Backend, error) { - return nil, errors.New("no way to reach that worker") + return nil, ErrNoWorkerDialer +} + +// deadDialFactory builds clients normally and fails the DIAL, which is the +// RUNNING half and by far the likelier one: the factory only fails when the +// wiring is absent, while the dial fails whenever the replica holding a +// worker's tunnel is momentarily unreachable, which one frontend restart +// produces for every worker that replica holds. +func deadDialFactory(cause error) BackendClientFactory { + GinkgoHelper() + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + return func(context.Context, string) (net.Conn, error) { return nil, cause } + }) + Expect(err).ToNot(HaveOccurred()) + return f } var _ = Describe("routing when the worker cannot be reached at all", func() { @@ -68,3 +86,99 @@ var _ = Describe("routing when the worker cannot be reached at all", func() { Expect(reg.decrementCalls).To(ContainElement("X:m")) }) }) + +var _ = Describe("routing when the worker's tunnel dial fails", func() { + // The reviewer's spec. It is the boundary test: the factory succeeds, the + // gRPC client is built, and the DIAL fails underneath with a + // cluster.ErrNoRoute. gRPC flattens that into codes.Unavailable, which is + // also what a dead backend produces, so without a way to carry the + // distinction past the package boundary a peer link blip is read as a dead + // process and the replica row is deleted after ONE miss. + loadedReg := func() *fakeModelRouter { + node := &BackendNode{ID: "X", Name: "node-x", Address: "10.0.0.1:50051"} + nm := &NodeModel{NodeID: "X", ModelName: "m", Address: "10.0.0.1:9001"} + return &fakeModelRouter{ + findAndLockNode: node, + findAndLockNM: nm, + loadedReplicaStatsByName: map[string][]ReplicaCandidate{"m": {{NodeID: "X", InFlight: 0}}}, + } + } + + route := func(reg *fakeModelRouter, cause error) error { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: &fakeUnloader{}, + ClientFactory: deadDialFactory(cause), + }) + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) + return err + } + + It("never reaps a replica whose OWNER replica is unreachable", func() { + reg := loadedReg() + Expect(route(reg, fmt.Errorf("through replica %q: %w: %w", "peer-2", cluster.ErrNoRoute, cluster.ErrPeerUnreachable))).To(HaveOccurred()) + Expect(reg.removeCalls).To(BeEmpty(), + "a worker whose OWNER replica is unreachable must never have its loaded models reaped") + }) + + It("never reaps a replica that has not dialled its tunnel yet", func() { + // The rolling-upgrade case end to end. A frontend-first upgrade puts + // every not-yet-restarted worker here at once, and every one of them is + // heartbeating and serving while it happens. + reg := loadedReg() + Expect(route(reg, fmt.Errorf("reaching node %q: %w", "X", cluster.ErrNoRoute))).To(HaveOccurred()) + Expect(reg.removeCalls).To(BeEmpty()) + }) + + It("still releases the reservation it took", func() { + reg := loadedReg() + Expect(route(reg, fmt.Errorf("%w", cluster.ErrNoRoute))).To(HaveOccurred()) + Expect(reg.decrementCalls).To(ContainElement("X:m")) + }) + + It("carries the cluster condition all the way across the package boundary", func() { + // Not just "something failed": the specific reason survives gRPC, which + // is what makes the five conditions usable on this side. If this ever + // reduces to a bare code, the consumers above are guessing again. + f := deadDialFactory(fmt.Errorf("through replica %q: %w: %w", "peer-2", cluster.ErrNoRoute, cluster.ErrPeerUnreachable)) + client, err := f.NewClientForNode("X", "10.0.0.1:9001", false) + Expect(err).ToNot(HaveOccurred()) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, _ = client.HealthCheck(ctx) + + unreached := unroutable(client) + Expect(unreached).To(MatchError(ErrWorkerUnroutable)) + Expect(unreached).To(MatchError(cluster.ErrNoRoute)) + Expect(unreached).To(MatchError(cluster.ErrPeerUnreachable)) + // And never absence, at either end of the trip. + Expect(unreached).ToNot(MatchError(cluster.ErrNoConnection)) + Expect(unreached).ToNot(MatchError(cluster.ErrInstanceNotFound)) + }) + + It("reports nothing for a client whose dial succeeded", func() { + // The other direction, so the seam cannot pass by always saying yes: a + // backend that genuinely died must still be reapable. + listener, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = listener.Close() }) + + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + var d net.Dialer + return func(ctx context.Context, _ string) (net.Conn, error) { + return d.DialContext(ctx, "tcp", listener.Addr().String()) + } + }) + Expect(err).ToNot(HaveOccurred()) + client, err := f.NewClientForNode("X", "10.0.0.1:9001", false) + Expect(err).ToNot(HaveOccurred()) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + // The listener accepts and speaks no gRPC, so the RPC fails while the + // DIAL succeeds. That is exactly a dead-ish backend on a reachable + // worker, and it must not read as unroutable. + _, _ = client.HealthCheck(ctx) + Expect(unroutable(client)).To(BeNil()) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index d73cdf6b6cb6..4f10775e1557 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -174,7 +174,16 @@ Four outcomes are kept apart on purpose, because they call for different actions Only the first is absence. The others are never reported as it, and that is not a stylistic preference: a scheduler told that a connected worker has gone away reclaims every model it is running. -Set `LOCALAI_WORKER_TUNNEL=false` on a worker to turn the tunnel off and go back to the frontend dialling the worker's advertised addresses. +#### There is no frontend-side fallback, and upgrade order matters + +`LOCALAI_WORKER_TUNNEL=false` still stops a worker dialling its tunnel, but it no longer has a frontend counterpart: after this change **no frontend path dials a worker's advertised address**, so a worker with the tunnel off is a worker the frontend cannot reach. Setting it is not a rollback. The rollback is to run the previous frontend release. + +That makes upgrade order matter, in one direction only: + +- **Upgrade the workers first, then the frontends.** A worker on the new build dials its tunnel and is reachable by frontends of either version, because the old frontend still dials its advertised address and the worker still listens. +- **Upgrading the frontends first** leaves every not-yet-restarted worker unroutable until it restarts. Those workers keep running their models and keep heartbeating, and the frontend reports them as unroutable rather than as gone: their `node_models` rows are left alone, nothing is rescheduled, and requests for those models fail loudly with "no route" until the worker reconnects. It is a degraded window, not an eviction, but it is a window, and doing it the other way round has none. + +A worker that cannot reach its frontend retries with exponential backoff and never gives up, so restarting a worker is all that is needed to close the window. ### The model load deadline scales with the checkpoint @@ -399,7 +408,7 @@ local-ai worker \ | `--registration-require-auth` | `LOCALAI_REGISTRATION_REQUIRE_AUTH` | `false` | Refuse to start the HTTP file-transfer server when no registration token is set (it would otherwise fail open) | | `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | Umbrella switch implying both `--registration-require-auth` and `--nats-require-auth` | | `--heartbeat-interval` | `LOCALAI_HEARTBEAT_INTERVAL` | `10s` | Interval between heartbeat pings | -| `--worker-tunnel` | `LOCALAI_WORKER_TUNNEL` | `true` | Hold one outbound multiplexed tunnel to the frontend and serve its requests over it, so this worker needs no inbound port (see [Worker tunnels](#worker-tunnels)) | +| `--worker-tunnel` | `LOCALAI_WORKER_TUNNEL` | `true` | Hold one outbound multiplexed tunnel to the frontend and serve its requests over it, so this worker needs no inbound port (see [Worker tunnels](#worker-tunnels)). Turning it off makes the worker unreachable: the frontend has no path that dials a worker's advertised address. | | `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS URL for backend installation and file staging | | `--nats-jwt` | `LOCALAI_NATS_JWT` | *(empty)* | Optional override for the `nats_jwt` returned at registration | | `--nats-user-seed` | `LOCALAI_NATS_USER_SEED` | *(empty)* | Optional override for `nats_user_seed` from registration | diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index e90bbb28fc5f..2bb013de075d 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -49,10 +49,29 @@ func NewClientWithDialer(address string, parallel bool, wd WatchDog, enableWatch // the address directly, which is the exact bypass this constructor exists // to close. c := buildClient(address, parallel, wd, enableWatchDog, token) - c.dialer = dialer + // Wrapped rather than stored bare, so every dial outcome is recorded. This + // is the seam that carries the reason a dial failed past gRPC, which + // flattens it into codes.Unavailable; see (*Client).LastDialError. + 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 that transport, rather than +// the process, is what failed. +// +// It is a separate interface and NOT part of Backend on purpose: only the +// handful of callers that act on the difference need it, and widening Backend +// would make every wrapper and every test double implement a method they have +// no answer for. +type DialErrorReporter interface { + LastDialError() error +} + 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 6131d3801bc7..e132086a013a 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -40,6 +40,14 @@ type Client struct { // keeps gRPC's own TCP dial, which is what every non-distributed caller // wants. dialer func(ctx context.Context, addr string) (net.Conn, error) + + // dialErrMu guards lastDialErr. Its own mutex rather than the embedded one: + // the embedded Mutex guards inFlight and is taken on every call, and a + // dialer runs underneath gRPC's own machinery where reentering it is not + // something this type can reason about. + dialErrMu sync.Mutex + lastDialErr error + sync.Mutex opMutex sync.Mutex wd WatchDog @@ -1422,3 +1430,40 @@ 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 from the most recent attempt by this client's +// custom dialer, or nil when the last attempt succeeded or there is no custom +// dialer. +// +// It exists because gRPC destroys the distinction its callers need. A dialer +// failure reaches an RPC as codes.Unavailable with the cause flattened into a +// message string, and codes.Unavailable is ALSO what a backend process that +// died produces. Those two call for opposite actions: a dead backend's registry +// row should be reaped, and a transport that could not reach a live backend +// must never cause one to be. Recording the error here is what lets a caller +// tell them apart, with the original error VALUE intact, so +// core/services/cluster's sentinels survive the trip. +// +// Scope, stated exactly. This is the last dial on this CLIENT, not the last +// dial for a particular RPC. A client used for one probe and closed gives exact +// attribution, which is how every reaping path in core/services/nodes uses it. +// A client shared across concurrent RPCs can attribute a dial failure to the +// wrong one; both directions of that error are safe, because a caller consults +// this only when its RPC already failed, and the outcomes are "treat a dead +// backend as unreachable-for-now" (the row survives one extra round) or "treat +// a transport failure as a backend failure" (the behaviour before this +// existed). +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 previous +// failure rather than leaving it, so a client that recovered does not keep +// reporting a dial error that no longer describes anything. +func (c *Client) recordDialErr(err error) { + c.dialErrMu.Lock() + c.lastDialErr = err + c.dialErrMu.Unlock() +} diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 322b11e36c96..1193a18d8c48 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" @@ -685,6 +686,20 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { // 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. + // + // "The node is down" is exactly what this can no longer conclude on + // its own. In distributed mode the client reaches the backend over + // the worker's tunnel, and a failure of THAT transport (the replica + // holding the tunnel is restarting, the worker has not dialled in + // yet after a frontend-first upgrade) arrives as the same + // codes.Unavailable a dead worker produces. Evicting on it would + // unload a model that is loaded and serving. The client records + // which of the two happened; see grpc.DialErrorReporter. + if dialErr := transportFailure(client); dialErr != nil { + xlog.Warn("Remote model health check could not reach the worker, keeping cached", + "model", s, "error", dialErr) + 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 { @@ -709,3 +724,20 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { m.MarkHealthy() return m } + +// transportFailure reports why a call never reached the backend, or nil when it +// did reach one. +// +// It is the one question that separates "this backend is gone" from "this +// process cannot currently get to it", and gRPC does not answer it: a dialer +// failure and a dead listener both surface as codes.Unavailable. A client with +// no custom transport answers nil, which is right for every locally spawned +// backend, where the address IS a socket on this machine and a failed +// connection really does mean the process died. +func transportFailure(client grpc.Backend) error { + reporter, ok := client.(grpc.DialErrorReporter) + if !ok { + return nil + } + return reporter.LastDialError() +} diff --git a/pkg/model/remote_unroutable_internal_test.go b/pkg/model/remote_unroutable_internal_test.go new file mode 100644 index 000000000000..25a440d800be --- /dev/null +++ b/pkg/model/remote_unroutable_internal_test.go @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: MIT + +package model + +import ( + "context" + "errors" + "net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" + "github.com/mudler/LocalAI/pkg/system" +) + +var _ = Describe("the health check on a remote model whose transport failed", func() { + // The fourth site of the same shape as the reconciler, the health monitor + // and the router, found by sweeping rather than by being named. + // + // checkIsLoaded evicts a remote model on a "connection error", which used + // to mean exactly one thing: the worker's socket did not answer. In + // distributed mode the client reaches the backend over the worker's tunnel, + // and a failure of THAT transport arrives as the same codes.Unavailable. + // Evicting on it unloads a model that is loaded and serving, on a worker + // that is heartbeating. + var ml *ModelLoader + + BeforeEach(func() { + systemState, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir())) + Expect(err).ToNot(HaveOccurred()) + ml = NewModelLoader(systemState) + }) + + It("keeps the model when the tunnel dial failed", func() { + client := grpc.NewClientWithDialer("10.0.0.1:9001", false, nil, false, "", + func(context.Context, string) (net.Conn, error) { + return nil, errors.New("cluster: no route from this replica to that worker") + }) + m := NewModelWithClient("remote-model", "10.0.0.1:9001", client) + ml.store.Set("remote-model", m) + + Expect(ml.checkIsLoaded("remote-model")).To(BeIdenticalTo(m), + "a model on a worker this frontend cannot route to must stay cached, not be unloaded") + _, stillThere := ml.store.Get("remote-model") + Expect(stillThere).To(BeTrue()) + }) + + It("still evicts a remote model whose worker WAS reached and did not answer", func() { + // The other direction, so the new check cannot pass by never evicting. + // No custom dialer, so the transport reports nothing and a connection + // error means what it always meant. + client := grpc.NewClientWithToken("127.0.0.1:1", false, nil, false, "") + m := NewModelWithClient("dead-model", "127.0.0.1:1", client) + ml.store.Set("dead-model", m) + + Expect(ml.checkIsLoaded("dead-model")).To(BeNil()) + _, stillThere := ml.store.Get("dead-model") + Expect(stillThere).To(BeFalse()) + }) +}) From b4d8e23abb9d4d5ba60e910237d44aacc18ae678 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 16:33:59 +0000 Subject: [PATCH 33/42] fix(grpc): let the transport answer through the wrappers, not only past gRPC Re-review round 2. One blocking defect, and it was the concern I filed myself last round and mis-scoped as a future trap. It was live, and it sat on the most destructive reaping path of the five. RouteResult.Client is an InFlightTrackingClient, over a FileStagingClient when a stager is configured. model_router puts that on the cached remote model and pkg/model's checkIsLoaded asks IT whether the transport failed. Both wrappers embed grpc.Backend, which does not declare LastDialError, so the type assertion read nil and the guard added last round fell straight through to the old eviction. That eviction sends backend.stop over NATS to every node holding the model and deletes every replica row, where the other sites delete one. The spec covering it built a bare client by hand, which is why it passed while production did not. This is the third time in this task a correct fix was disarmed one layer out, so the fix is a mechanism rather than two methods. BackendUnwrapper is one line per decorator, LastDialErrorOf walks the chain, and both consumers now call it instead of each keeping its own assertion. One implementation, no per-caller policy to get wrong. Sweeping every type that embeds or holds a grpc.Backend found a third decorator the review had not named, and it is itself a reaping consumer of the same collapsed signal. ConnectionEvictingClient is built for remote models in initializers.go and its evict callback runs ShutdownModel; it fires during INFERENCE rather than on a health check, so a tunnel blip mid-request was enough to stop a model that was loaded and serving. It consults the transport first now. A locally spawned backend has no custom transport, so that path is unchanged byte for byte. Everything else touching a Backend is a consumer rather than a decorator; there is no fourth. The probe cache joiner shape is pinned. It was the right design last round with nothing holding it: the mutation back to a closed-over variable passed all 602 specs in the package. Eight goroutines coalesced on a probe that blocks on a channel now assert every joiner gets the leader's REASON and not just its answer, which is the difference between a leader declining to reap and its seven joiners reaping on the leader's own observation. The LastDialError scope note claimed an exactness it does not have at checkIsLoaded, which reads a shared long-lived client after releasing opMutex. It now says which caller is not exact, why the imprecision is accepted there, and what making it exact would cost. The four-outcome table in the docs still said a worker with no live owner is treated as absent and rescheduled, contradicting the code and the paragraph nine lines below it. None of those outcomes is absence any more, and the table says so, names the fifth, and points at the heartbeat as the thing that does decide presence. Five mutations, each reddening named specs, including the two the reviewer found surviving. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/nodes/file_staging_client.go | 8 ++ core/services/nodes/inflight.go | 23 +++- core/services/nodes/interfaces.go | 11 +- core/services/nodes/probe_cache_test.go | 55 +++++++++ core/services/nodes/wrapper_transport_test.go | 115 ++++++++++++++++++ docs/content/features/distributed-mode.md | 9 +- pkg/grpc/backend.go | 48 ++++++++ pkg/grpc/client.go | 29 +++-- pkg/model/connection_evicting_client.go | 33 ++++- pkg/model/loader.go | 17 ++- pkg/model/remote_unroutable_internal_test.go | 47 +++++++ 11 files changed, 362 insertions(+), 33 deletions(-) create mode 100644 core/services/nodes/wrapper_transport_test.go diff --git a/core/services/nodes/file_staging_client.go b/core/services/nodes/file_staging_client.go index bfc202c8205d..d73a59267c4e 100644 --- a/core/services/nodes/file_staging_client.go +++ b/core/services/nodes/file_staging_client.go @@ -37,6 +37,14 @@ type FileStagingClient struct { remoteModelPath string // set during LoadModel from staged ModelPath } +// Unwrap exposes the client this one decorates, so grpc.LastDialErrorOf can see +// past it. Without it a staged client answers "the transport was fine" for +// every dial, because embedding grpc.Backend inherits only what Backend +// declares and DialErrorReporter is deliberately not on Backend. +func (f *FileStagingClient) Unwrap() grpc.Backend { return f.Backend } + +var _ grpc.BackendUnwrapper = (*FileStagingClient)(nil) + // NewFileStagingClient creates a new file staging wrapper. func NewFileStagingClient(inner grpc.Backend, stager FileStager, nodeID string) *FileStagingClient { return &FileStagingClient{ diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index 3102a3254804..f46bb3e5b6cc 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -30,10 +30,15 @@ import ( type InFlightTrackingClient struct { grpc.ControlBackend // passthrough for control-plane / streaming-constructor methods inner grpc.InferenceBackend // tracked inference methods delegate here - registry InFlightTracker - nodeID string - modelName string - replicaIndex int + // wrapped is the SAME object as ControlBackend and inner, kept at its full + // type so Unwrap can hand it back. The two fields above are deliberately + // narrowed to the sub-interfaces, which is what gives the compile-time + // guarantee below, and neither of them can be returned as a grpc.Backend. + wrapped grpc.Backend + registry InFlightTracker + nodeID string + modelName string + replicaIndex int firstOnce sync.Once // guards onFirstComplete onFirstComplete func() // called once after the first tracked inference call completes @@ -44,11 +49,21 @@ type InFlightTrackingClient struct { // InferenceBackend method is left unwrapped. var _ grpc.Backend = (*InFlightTrackingClient)(nil) +// And it must stay transparent to grpc.LastDialErrorOf. This is the wrapper +// SmartRouter puts on every routed client, so a remote model's cached client is +// one of these; without Unwrap, the transport guard in pkg/model reads nil for +// every model the router produced and evicts on a tunnel blip. +var _ grpc.BackendUnwrapper = (*InFlightTrackingClient)(nil) + +// Unwrap exposes the client this one decorates. +func (c *InFlightTrackingClient) Unwrap() grpc.Backend { return c.wrapped } + // NewInFlightTrackingClient wraps a gRPC backend client with in-flight tracking. func NewInFlightTrackingClient(inner grpc.Backend, registry InFlightTracker, nodeID, modelName string, replicaIndex int) *InFlightTrackingClient { return &InFlightTrackingClient{ ControlBackend: inner, inner: inner, + wrapped: inner, registry: registry, nodeID: nodeID, modelName: modelName, diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index d1028983843b..c8aa23323db3 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -208,13 +208,12 @@ var ErrNoWorkerDialer = fmt.Errorf("%w: no worker tunnel dialer is configured", // // A client that reports nothing (no custom dialer, or a test double) yields // nil, which means "the call reached a backend" and preserves the behaviour -// every non-distributed caller has always had. +// every non-distributed caller has always had. Decorators are looked through; +// see grpc.BackendUnwrapper for why that is not optional. func unroutable(client grpc.Backend) error { - reporter, ok := client.(grpc.DialErrorReporter) - if !ok { - return nil - } - dialErr := reporter.LastDialError() + // LastDialErrorOf and not a type assertion: the assertion could not see + // past a decorator, and SmartRouter hands every routed client out wrapped. + dialErr := grpc.LastDialErrorOf(client) if dialErr == nil { return nil } diff --git a/core/services/nodes/probe_cache_test.go b/core/services/nodes/probe_cache_test.go index 58e6fa111cb9..d28c4c7b7ade 100644 --- a/core/services/nodes/probe_cache_test.go +++ b/core/services/nodes/probe_cache_test.go @@ -1,6 +1,7 @@ package nodes import ( + "errors" "sync" "sync/atomic" "time" @@ -104,6 +105,60 @@ var _ = Describe("probeCache", func() { } }) + It("hands every coalesced joiner the leader's REASON, not just its answer", func() { + // The hole this shape exists to close, and the one a closed-over + // variable reintroduces. The reason is written only in the goroutine + // that runs the probe; every caller coalesced into that flight would + // read its own unset variable and see nil. In production that means the + // leader correctly declines to reap a replica on an unreachable worker + // while all seven joiners reap it, on the leader's own observation. + c := newProbeCache(time.Minute) + unreached := errors.New("no route to the worker") + + // The probe blocks until every goroutine is inside flight.Do, so the + // joiners are genuinely coalesced rather than serialised. Released by a + // channel, so nothing here waits on a clock. + entered := make(chan struct{}) + release := make(chan struct{}) + var calls int32 + probe := func() (bool, error) { + atomic.AddInt32(&calls, 1) + close(entered) + <-release + return false, unreached + } + + const N = 8 + start := make(chan struct{}) + var wg sync.WaitGroup + reasons := make([]error, N) + alive := make([]bool, N) + for i := 0; i < N; i++ { + wg.Add(1) + go func(i int) { + defer GinkgoRecover() + defer wg.Done() + <-start + alive[i], reasons[i] = c.DoOrCachedResult("k", probe) + }(i) + } + close(start) + + // Only unblock the leader once at least one goroutine is inside the + // probe; the rest are then either waiting on the flight or about to be. + <-entered + close(release) + wg.Wait() + + Expect(atomic.LoadInt32(&calls)).To(Equal(int32(1)), + "singleflight must collapse %d concurrent probes into one", N) + for i := range reasons { + Expect(alive[i]).To(BeFalse(), "goroutine %d saw a different answer", i) + Expect(reasons[i]).To(MatchError(unreached), + "goroutine %d joined the flight and got the answer without the reason, which is how a joiner reaps what the leader would not", i) + } + }) + It("treats different keys independently", func() { c := newProbeCache(time.Minute) var aCalls, bCalls int32 diff --git a/core/services/nodes/wrapper_transport_test.go b/core/services/nodes/wrapper_transport_test.go new file mode 100644 index 000000000000..ff504929a60c --- /dev/null +++ b/core/services/nodes/wrapper_transport_test.go @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: MIT + +package nodes + +import ( + "context" + "errors" + "net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" + grpc "github.com/mudler/LocalAI/pkg/grpc" +) + +// The reviewer's spec, plus the production shape it was pointing at. +// +// The guard added at the fourth reap site asks the client whether the TRANSPORT +// failed. In production that client is not the one the factory built: SmartRouter +// hands out result.Client, which is an *InFlightTrackingClient, over a +// *FileStagingClient whenever a stager is configured. Both embed grpc.Backend, +// which does not declare LastDialError, so a type assertion on the outermost +// type read nil and the guard was inert for exactly the models the router +// produces. Every spec that constructed a raw client by hand passed anyway. +// +// This is the third time in this task that a correct fix was disarmed by a +// layer further out, which is why the mechanism is now one walker rather than a +// per-caller assertion. +var _ = Describe("the transport answer through the wrappers the router builds", func() { + var ( + cause error + raw grpc.Backend + ) + + BeforeEach(func() { + cause = errors.New("cluster: no route from this replica to that worker") + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + return func(context.Context, string) (net.Conn, error) { return nil, cause } + }) + Expect(err).ToNot(HaveOccurred()) + raw, err = f.NewClientForNode("X", "10.0.0.1:9001", false) + Expect(err).ToNot(HaveOccurred()) + + // Provoke one dial so there is something to report. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, _ = raw.HealthCheck(ctx) + }) + + It("the raw factory client reports, as designed", func() { + Expect(unroutable(raw)).To(MatchError(ErrWorkerUnroutable)) + }) + + It("reports through the in-flight tracker, which is what RouteResult.Client is", func() { + tracked := NewInFlightTrackingClient(raw, &fakeModelRouter{}, "X", "m", 0) + Expect(unroutable(tracked)).To(MatchError(ErrWorkerUnroutable)) + }) + + It("reports through the file staging client, which buildClientForAddr adds", func() { + staged := NewFileStagingClient(raw, nil, "X") + Expect(unroutable(staged)).To(MatchError(ErrWorkerUnroutable)) + }) + + It("reports through BOTH, nested the way production nests them", func() { + // buildClientForAddr wraps in staging, newRouteResult wraps that in + // tracking, model_router puts the result on the cached model, and + // pkg/model's checkIsLoaded asks it. Two layers, and a walker that + // stopped at one would still be wrong here. + nested := NewInFlightTrackingClient(NewFileStagingClient(raw, nil, "X"), &fakeModelRouter{}, "X", "m", 0) + Expect(unroutable(nested)).To(MatchError(ErrWorkerUnroutable)) + }) + + It("still reports nothing through the wrappers when the dial succeeded", func() { + // The other direction, so forwarding cannot pass by always answering + // "unroutable": a backend that genuinely died must still be reapable + // through the same wrappers. + listener, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = listener.Close() }) + + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + var d net.Dialer + return func(ctx context.Context, _ string) (net.Conn, error) { + return d.DialContext(ctx, "tcp", listener.Addr().String()) + } + }) + Expect(err).ToNot(HaveOccurred()) + live, err := f.NewClientForNode("X", "10.0.0.1:9001", false) + Expect(err).ToNot(HaveOccurred()) + _, _ = live.HealthCheck(context.Background()) + + nested := NewInFlightTrackingClient(NewFileStagingClient(live, nil, "X"), &fakeModelRouter{}, "X", "m", 0) + Expect(unroutable(nested)).To(BeNil()) + }) + + It("keeps the cluster condition matchable through the wrappers", func() { + // Not merely "something failed". The five conditions have to survive + // the decorators as well as gRPC, or the consumers are guessing again. + routed := errors.New("x") + f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) { + return func(context.Context, string) (net.Conn, error) { return nil, routed } + }) + Expect(err).ToNot(HaveOccurred()) + c, err := f.NewClientForNode("X", "10.0.0.1:9001", false) + Expect(err).ToNot(HaveOccurred()) + routed = cluster.ErrNoRoute + _, _ = c.HealthCheck(context.Background()) + + nested := NewInFlightTrackingClient(NewFileStagingClient(c, nil, "X"), &fakeModelRouter{}, "X", "m", 0) + got := unroutable(nested) + Expect(got).To(MatchError(cluster.ErrNoRoute)) + Expect(got).ToNot(MatchError(cluster.ErrNoConnection)) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 4f10775e1557..a4700ec4f3e1 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -163,16 +163,19 @@ A worker's tunnel lands on exactly one replica, so with N replicas behind a load The dialling replica states how much time its own client has left in the frame that opens the relayed stream, and the owner bounds its work by the smaller of that and its own 15s ceiling. Neither number can lengthen the other: a patient client cannot park the owning replica, and an impatient one cannot be kept waiting on a budget it did not ask for. -Four outcomes are kept apart on purpose, because they call for different actions: +These outcomes are kept apart on purpose, because they call for different actions: | Outcome | What it means | What acts on it | |---|---|---| -| No live owner | No replica holds this worker's tunnel | The worker is treated as absent; its models can be rescheduled | +| No live owner | No replica holds this worker's tunnel | No route right now; the worker's models are **left alone** | | Not the owner | The routing was stale | Resolve the owner again | | Peer unreachable | A replica exists and will not answer | Retry | +| No relay path | This replica cannot reach the owner at all | Report; requests here fail until it can | | The worker refused | The worker answered and said no | Report; the worker is connected | -Only the first is absence. The others are never reported as it, and that is not a stylistic preference: a scheduler told that a connected worker has gone away reclaims every model it is running. +**None of them is absence.** A worker's presence is its **heartbeat**, and a route to it is a separate fact that can be false while the worker is registered, heartbeating and serving every request another replica sends it. So the frontend answers "no route", never "this worker is gone", and nothing on this list causes a model to be rescheduled or a `node_models` row to be deleted. + +That distinction is the whole point rather than a nicety. A scheduler told that a connected worker has gone away stops its backend and reclaims every model it is running, and the events that produce "no route" are ordinary ones: a frontend replica restarting, an ownership row a moment stale, a worker that has not dialled its tunnel yet. A worker is treated as absent only when its **heartbeat** goes stale, which is a separate mechanism with its own threshold (see `--stale-node-threshold`). #### There is no frontend-side fallback, and upgrade order matters diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 2bb013de075d..1126e39b7269 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -72,6 +72,54 @@ type DialErrorReporter interface { LastDialError() error } +// BackendUnwrapper is implemented by a Backend that DECORATES another one. +// +// Every wrapper in this codebase must implement it, and the reason is a defect +// that shipped: a wrapper embeds the Backend interface, so it inherits every +// declared method and NOTHING else. DialErrorReporter is deliberately not +// declared on Backend, so a wrapped client silently stopped answering "did the +// transport fail" and the guard built on that answer read nil in production +// while passing every spec that constructed a raw client by hand. +// +// Implementing this is what makes a decorator transparent to LastDialErrorOf, +// and it is one line rather than a re-implementation per wrapper, so there is +// no per-wrapper policy to get wrong. +type BackendUnwrapper interface { + Unwrap() Backend +} + +// maxBackendUnwrapDepth bounds the walk below. Three wrappers exist today and +// they nest at most two deep; the bound is a guard against a cycle a future +// wrapper could introduce, not a limit anything real approaches. +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. +// +// It is the single implementation of that question. Its callers +// (core/services/nodes and pkg/model) each had their own type assertion, and an +// assertion cannot see past a wrapper: in production the client handed to +// pkg/model is an *InFlightTrackingClient over a *FileStagingClient over the +// real one, so both callers were asking a wrapper that had no answer and +// reading nil as "the transport was fine". +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 +} + 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 e132086a013a..fe4a5d182131 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -1444,15 +1444,26 @@ func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts .. // tell them apart, with the original error VALUE intact, so // core/services/cluster's sentinels survive the trip. // -// Scope, stated exactly. This is the last dial on this CLIENT, not the last -// dial for a particular RPC. A client used for one probe and closed gives exact -// attribution, which is how every reaping path in core/services/nodes uses it. -// A client shared across concurrent RPCs can attribute a dial failure to the -// wrong one; both directions of that error are safe, because a caller consults -// this only when its RPC already failed, and the outcomes are "treat a dead -// backend as unreachable-for-now" (the row survives one extra round) or "treat -// a transport failure as a backend failure" (the behaviour before this -// existed). +// Scope, stated exactly, including where it is NOT exact. +// +// This is the last dial on this CLIENT, not the last dial for a particular RPC. +// Three of the four callers build a client for one probe and close it, so +// attribution there is exact. The fourth, pkg/model's checkIsLoaded, reads the +// model's long-lived SHARED client and consults this after HealthCheck has +// released opMutex, so a concurrent RPC on the same client can record or clear +// the value inside that window. An earlier version of this comment claimed +// exactness for all four; it was wrong. +// +// The imprecision is accepted there rather than designed away, and the reason +// is which way it can go. A caller consults this only when its own RPC already +// failed, so the two outcomes are: a concurrent dial FAILURE makes a genuinely +// dead backend look unreachable-for-now, and its row survives one extra round +// until the transport recovers; or a concurrent dial SUCCESS clears the value +// and a transport failure reads as a backend failure, which is exactly the +// behaviour that existed before any of this. Neither is a new hazard, and the +// second requires a transport that recovered inside the window. Making it exact +// would mean threading a per-call handle through every Backend method, which is +// a far larger change than the failure it would prevent. func (c *Client) LastDialError() error { c.dialErrMu.Lock() defer c.dialErrMu.Unlock() diff --git a/pkg/model/connection_evicting_client.go b/pkg/model/connection_evicting_client.go index 00d42d200f96..bde2333bfa61 100644 --- a/pkg/model/connection_evicting_client.go +++ b/pkg/model/connection_evicting_client.go @@ -22,6 +22,12 @@ type ConnectionEvictingClient struct { once sync.Once } +var _ grpc.BackendUnwrapper = (*ConnectionEvictingClient)(nil) + +// Unwrap exposes the client this one decorates, so grpc.LastDialErrorOf can see +// past it. +func (c *ConnectionEvictingClient) Unwrap() grpc.Backend { return c.Backend } + func newConnectionEvictingClient(inner grpc.Backend, modelID string, evict func()) grpc.Backend { return &ConnectionEvictingClient{ Backend: inner, @@ -31,13 +37,28 @@ func newConnectionEvictingClient(inner grpc.Backend, modelID string, evict func( } func (c *ConnectionEvictingClient) checkErr(err error) { - if err != nil && isConnectionError(err) { - c.once.Do(func() { - xlog.Warn("Connection error during inference, evicting model from cache", - "model", c.modelID, "error", err) - c.evict() - }) + if err == nil || !isConnectionError(err) { + return + } + // The fifth site of the same shape, and the one reached during INFERENCE + // rather than a health check. evict() runs ShutdownModel, which for a remote + // model sends backend.stop over NATS to every node holding it and deletes + // every replica row. In distributed mode the client underneath reaches the + // backend over the worker's tunnel, and a failure of THAT transport arrives + // as the same codes.Unavailable a dead backend produces; evicting on it + // stops a model that is loaded and serving, on a worker that is + // heartbeating. A locally spawned backend has no custom transport, so this + // reports nil and the behaviour there is exactly what it always was. + if dialErr := grpc.LastDialErrorOf(c.Backend); dialErr != nil { + xlog.Warn("Inference failed because the worker could not be reached; keeping the model", + "model", c.modelID, "error", dialErr) + return } + c.once.Do(func() { + xlog.Warn("Connection error during inference, evicting model from cache", + "model", c.modelID, "error", err) + c.evict() + }) } // --- Intercepted inference methods --- diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 1193a18d8c48..91fdbff672d2 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -695,6 +695,12 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { // codes.Unavailable a dead worker produces. Evicting on it would // unload a model that is loaded and serving. The client records // which of the two happened; see grpc.DialErrorReporter. + // The client here is long-lived and shared, so this reads the last + // dial on it rather than the one this HealthCheck made; see + // (*grpc.Client).LastDialError for why that imprecision is + // accepted. Both directions of it land on behaviour that already + // existed, and the common case (a worker with no route at all) has + // no concurrent success to clear the value. if dialErr := transportFailure(client); dialErr != nil { xlog.Warn("Remote model health check could not reach the worker, keeping cached", "model", s, "error", dialErr) @@ -735,9 +741,10 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { // backend, where the address IS a socket on this machine and a failed // connection really does mean the process died. func transportFailure(client grpc.Backend) error { - reporter, ok := client.(grpc.DialErrorReporter) - if !ok { - return nil - } - return reporter.LastDialError() + // LastDialErrorOf and not a type assertion. The client reaching this + // function for a routed remote model is an *InFlightTrackingClient, often + // over a *FileStagingClient, and an assertion on the outermost type reads + // nil for both: they embed grpc.Backend, which does not declare + // LastDialError. That is exactly how this guard shipped inert. + return grpc.LastDialErrorOf(client) } diff --git a/pkg/model/remote_unroutable_internal_test.go b/pkg/model/remote_unroutable_internal_test.go index 25a440d800be..06883176f796 100644 --- a/pkg/model/remote_unroutable_internal_test.go +++ b/pkg/model/remote_unroutable_internal_test.go @@ -11,6 +11,7 @@ import ( . "github.com/onsi/gomega" grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/system" ) @@ -59,3 +60,49 @@ var _ = Describe("the health check on a remote model whose transport failed", fu Expect(stillThere).To(BeFalse()) }) }) + +var _ = Describe("the eviction wrapper on a remote model whose transport failed", func() { + // The FIFTH site of the same shape, found by sweeping the decorators rather + // than being named. initializers.go builds this wrapper for exactly the + // remote models the router produces, and its evict callback runs + // ShutdownModel, which sends backend.stop over NATS to every node holding + // the model and deletes every replica row. It fires during INFERENCE, not + // on a health check, so a tunnel blip mid-request was enough. + failingDial := func(cause error) grpc.Backend { + return grpc.NewClientWithDialer("10.0.0.1:9001", false, nil, false, "", + func(context.Context, string) (net.Conn, error) { return nil, cause }) + } + + It("does not evict when the worker could not be reached", func() { + evicted := 0 + client := newConnectionEvictingClient( + failingDial(errors.New("cluster: no route from this replica to that worker")), + "remote-model", func() { evicted++ }) + + _, err := client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(err).To(HaveOccurred()) + Expect(evicted).To(BeZero(), + "a worker this frontend cannot route to must not have its backend stopped and its rows deleted") + }) + + It("still evicts when the worker WAS reached and the connection failed", func() { + // The other direction. No custom dialer, so nothing reports a transport + // failure and a connection error means what it always meant. + evicted := 0 + client := newConnectionEvictingClient( + grpc.NewClientWithToken("127.0.0.1:1", false, nil, false, ""), + "dead-model", func() { evicted++ }) + + _, err := client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(err).To(HaveOccurred()) + Expect(evicted).To(Equal(1)) + }) + + It("is transparent to the transport question, so a wrapper of it still sees through", func() { + client := newConnectionEvictingClient( + failingDial(errors.New("cluster: no route from this replica to that worker")), + "remote-model", func() {}) + _, _ = client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(grpc.LastDialErrorOf(client)).ToNot(BeNil()) + }) +}) From b8d47cc29bfee3751df37bd8c157762f622a03f5 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 17:46:50 +0000 Subject: [PATCH 34/42] test(nodes): make the joiner spec deterministic, and gate the wrapper shape in lint Re-review round 2. One blocking item, and it was a spec I wrote: eight goroutines raced at the probe cache and nothing made them coalesce, so a straggler that missed the flight re-entered the probe and double-closed a channel. It panicked about one run in three and took the four-suite race block down. The green verification I reported was not reproducible, which means one green run was never evidence for a spec that coordinates goroutines. Its comment claimed the probe blocked until every goroutine was inside flight.Do, and that gap was exactly the panic: the comment described the design intended rather than the one written. It is deterministic now rather than tolerant. singleflight.DoChan registers its channel on an in-flight call under the group's own mutex and returns without running its function, so calling it while the leader is provably parked inside the probe joins that exact flight with no window and no dependence on the scheduler. The spec asserts the join really happened, that the joiner got the reason and not only the answer, and that the probe ran once; the entered channel is sent on rather than closed so a second probe fails an assertion instead of panicking. Twenty runs green under race against the committed code, five out of five red on the mutation back to a closed-over variable. The future-decorator gap is closed in the lint gate, but not the way the review suggested, and the reason is worth recording. HasMethod rejects inline signatures outright, its method-reference form needs a package ruleguard's own typechecker can import and that typechecker cannot import this module, and Implements tests the value method set while every Unwrap is on a pointer receiver, so it fired on all three wrappers that already had one. So the safe shape is structural instead. grpc.WrappedBackend gives the same pass-through method set plus Unwrap on a value receiver, and a decorator that embeds it is transparent by construction; forgetting stops being expressible rather than merely discouraged, which is the move loopbackService already makes in the worker. FileStagingClient and ConnectionEvictingClient embed it and their hand-written Unwrap methods are gone. The ruleguard rule then only has to catch the raw embedding, needs no type filter, and cannot misfire. It was verified to fire on a throwaway wrapper and stay silent on a correct one, and reports nothing across core and pkg with the baseline disabled. InFlightTrackingClient is the one exception and says why in a nolint: it embeds ControlBackend deliberately so that leaving an inference method unwrapped breaks the build, and WrappedBackend embeds the full interface, so adopting it would silently restore pass-through for every inference method and delete that guarantee. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- .golangci.yml | 22 ++++++ core/services/nodes/file_staging_client.go | 18 ++--- core/services/nodes/inflight.go | 10 +++ core/services/nodes/probe_cache_test.go | 83 ++++++++++++++-------- go.mod | 1 + go.sum | 2 + hack/lint/backend_wrappers.go | 46 ++++++++++++ pkg/grpc/backend.go | 24 +++++++ pkg/model/connection_evicting_client.go | 12 ++-- 9 files changed, 169 insertions(+), 49 deletions(-) create mode 100644 hack/lint/backend_wrappers.go diff --git a/.golangci.yml b/.golangci.yml index d25d1ccb4789..7a68e69a75f2 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -19,7 +19,22 @@ linters: - staticcheck enable: - forbidigo + # gocritic is enabled for ONE checker: ruleguard, which runs the rules in + # hack/lint/. Every other gocritic check is off (disable-all below), so + # this adds no style noise; it is here purely as the gate that catches a + # gRPC backend wrapper written without Unwrap. See + # hack/lint/backend_wrappers.go for why that cannot be a compile-time + # assertion. + - gocritic settings: + gocritic: + disable-all: true + enabled-checks: + - ruleguard + settings: + ruleguard: + failOn: all + rules: '${base-path}/hack/lint/backend_wrappers.go' forbidigo: forbid: - pattern: '^t\.Errorf$' @@ -126,3 +141,10 @@ linters: - path: ^backend/go/whisper/sources/ text: 'http\.(DefaultClient|Get|Post|PostForm|Head)' linters: [forbidigo] + # Test doubles embed grpc.Backend to inherit the interface's method set + # over a NIL value; they decorate nothing, hold no inner client, and have + # no transport answer to forward. The rule targets production wrappers, + # which is where swallowing that answer deletes replica rows. + # gocritic here is only the backend-wrapper ruleguard rule. + - path: _test\.go$ + linters: [gocritic] diff --git a/core/services/nodes/file_staging_client.go b/core/services/nodes/file_staging_client.go index d73a59267c4e..2fdedc3d219b 100644 --- a/core/services/nodes/file_staging_client.go +++ b/core/services/nodes/file_staging_client.go @@ -29,28 +29,22 @@ import ( // Methods that require no file staging are inherited from the embedded // grpc.Backend; only methods with staging logic are overridden below. type FileStagingClient struct { - grpc.Backend // embedded for pass-through of non-staging methods - stager FileStager - nodeID string + grpc.WrappedBackend // pass-through of non-staging methods, plus Unwrap + stager FileStager + nodeID string mu sync.RWMutex remoteModelPath string // set during LoadModel from staged ModelPath } -// Unwrap exposes the client this one decorates, so grpc.LastDialErrorOf can see -// past it. Without it a staged client answers "the transport was fine" for -// every dial, because embedding grpc.Backend inherits only what Backend -// declares and DialErrorReporter is deliberately not on Backend. -func (f *FileStagingClient) Unwrap() grpc.Backend { return f.Backend } - var _ grpc.BackendUnwrapper = (*FileStagingClient)(nil) // NewFileStagingClient creates a new file staging wrapper. func NewFileStagingClient(inner grpc.Backend, stager FileStager, nodeID string) *FileStagingClient { return &FileStagingClient{ - Backend: inner, - stager: stager, - nodeID: nodeID, + WrappedBackend: grpc.WrappedBackend{Backend: inner}, + stager: stager, + nodeID: nodeID, } } diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index f46bb3e5b6cc..8bf98584a821 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -27,6 +27,16 @@ import ( // interface therefore breaks this file's build (see the var assertion below) // until it is wrapped with track() - so a new inference path can't be added // without an in-flight accounting decision. +// The ruleguard rule wants grpc.WrappedBackend here, and this is the one +// decorator that cannot use it. Embedding ControlBackend rather than Backend is +// what produces the compile-time guarantee below: leave an InferenceBackend +// method unwrapped and this type stops satisfying grpc.Backend. WrappedBackend +// embeds the FULL interface, so adopting it would silently restore +// pass-through for every inference method and delete that guarantee. The +// transparency the rule protects is provided explicitly instead, by the Unwrap +// and the grpc.BackendUnwrapper assertion above. +// +//nolint:gocritic // embeds ControlBackend on purpose; Unwrap is declared explicitly type InFlightTrackingClient struct { grpc.ControlBackend // passthrough for control-plane / streaming-constructor methods inner grpc.InferenceBackend // tracked inference methods delegate here diff --git a/core/services/nodes/probe_cache_test.go b/core/services/nodes/probe_cache_test.go index d28c4c7b7ade..8d1e059127a8 100644 --- a/core/services/nodes/probe_cache_test.go +++ b/core/services/nodes/probe_cache_test.go @@ -8,6 +8,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "golang.org/x/sync/singleflight" ) var _ = Describe("probeCache", func() { @@ -112,51 +113,75 @@ var _ = Describe("probeCache", func() { // read its own unset variable and see nil. In production that means the // leader correctly declines to reap a replica on an unreachable worker // while all seven joiners reap it, on the leader's own observation. + // + // The FIRST version of this spec raced eight goroutines at the cache + // and hoped they coalesced. Nothing made them: a goroutine that arrived + // after the leader's flight finished started its own, re-entered the + // probe and double-closed a channel, so the spec panicked about one run + // in three. Its comment claimed the probe blocked until every goroutine + // was inside flight.Do, which was the design intended rather than the + // one written, and that gap was exactly the panic. + // + // This version does not hope. singleflight.DoChan registers its channel + // on the in-flight call under the group's own mutex and returns WITHOUT + // running its function (x/sync@v0.22.0 singleflight.go:127-132), so + // calling it while the leader is provably parked inside the probe joins + // that exact flight, with no window and no scheduler dependency. The + // group is reachable because this spec lives in the package. c := newProbeCache(time.Minute) unreached := errors.New("no route to the worker") - // The probe blocks until every goroutine is inside flight.Do, so the - // joiners are genuinely coalesced rather than serialised. Released by a - // channel, so nothing here waits on a clock. - entered := make(chan struct{}) + // Buffered, and sent on rather than closed: a probe that somehow ran + // twice must fail an assertion, not panic and take the suite with it. + entered := make(chan struct{}, 4) release := make(chan struct{}) var calls int32 probe := func() (bool, error) { atomic.AddInt32(&calls, 1) - close(entered) + entered <- struct{}{} <-release return false, unreached } - const N = 8 - start := make(chan struct{}) - var wg sync.WaitGroup - reasons := make([]error, N) - alive := make([]bool, N) - for i := 0; i < N; i++ { - wg.Add(1) - go func(i int) { - defer GinkgoRecover() - defer wg.Done() - <-start - alive[i], reasons[i] = c.DoOrCachedResult("k", probe) - }(i) + type leaderResult struct { + alive bool + unreached error } - close(start) + leader := make(chan leaderResult, 1) + go func() { + defer GinkgoRecover() + alive, reason := c.DoOrCachedResult("k", probe) + leader <- leaderResult{alive: alive, unreached: reason} + }() + + // The leader is now inside the probe, so the group holds an entry for + // "k" and will hold it until the probe returns. + Eventually(entered, "10s").Should(Receive()) + + // Deterministically coalesced. This function must never run; if the + // join failed it would, and the assertion below on the probe count + // would catch it too. + joined := c.flight.DoChan("k", func() (any, error) { + Fail("DoChan started its own flight, so nothing was coalesced") + return false, nil + }) - // Only unblock the leader once at least one goroutine is inside the - // probe; the rest are then either waiting on the flight or about to be. - <-entered close(release) - wg.Wait() + + var got leaderResult + Eventually(leader, "10s").Should(Receive(&got)) + Expect(got.alive).To(BeFalse()) + Expect(got.unreached).To(MatchError(unreached), "the caller that RAN the probe must get the reason") + + var shared singleflight.Result + Eventually(joined, "10s").Should(Receive(&shared)) + Expect(shared.Shared).To(BeTrue(), "this caller did not actually join the leader's flight") + Expect(shared.Val).To(Equal(false)) + Expect(shared.Err).To(MatchError(unreached), + "a joiner got the answer without the reason, which is how a joiner reaps what the leader would not") Expect(atomic.LoadInt32(&calls)).To(Equal(int32(1)), - "singleflight must collapse %d concurrent probes into one", N) - for i := range reasons { - Expect(alive[i]).To(BeFalse(), "goroutine %d saw a different answer", i) - Expect(reasons[i]).To(MatchError(unreached), - "goroutine %d joined the flight and got the answer without the reason, which is how a joiner reaps what the leader would not", i) - } + "the probe must have run exactly once") }) It("treats different keys independently", func() { diff --git a/go.mod b/go.mod index c3ded1cc3712..e1b1a72aabf1 100644 --- a/go.mod +++ b/go.mod @@ -54,6 +54,7 @@ require ( github.com/otiai10/openaigo v1.7.0 github.com/phayes/freeport v0.0.0-20220201140144-74d24b5ae9f5 github.com/prometheus/client_golang v1.23.2 + github.com/quasilyte/go-ruleguard/dsl v0.3.23 github.com/robfig/cron/v3 v3.0.1 github.com/russross/blackfriday v1.6.0 github.com/sashabaranov/go-openai v1.41.2 diff --git a/go.sum b/go.sum index 6a0be2c19cc8..a510df3dfd3b 100644 --- a/go.sum +++ b/go.sum @@ -1206,6 +1206,8 @@ github.com/prometheus/otlptranslator v1.0.0 h1:s0LJW/iN9dkIH+EnhiD3BlkkP5QVIUVEo github.com/prometheus/otlptranslator v1.0.0/go.mod h1:vRYWnXvI6aWGpsdY/mOT/cbeVRBlPWtBNDb7kGR3uKM= github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= +github.com/quasilyte/go-ruleguard/dsl v0.3.23 h1:lxjt5B6ZCiBeeNO8/oQsegE6fLeCzuMRoVWSkXC4uvY= +github.com/quasilyte/go-ruleguard/dsl v0.3.23/go.mod h1:KeCP03KrjuSO0H1kTuZQCWlQPulDV6YMIXmpQss17rU= github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw= diff --git a/hack/lint/backend_wrappers.go b/hack/lint/backend_wrappers.go new file mode 100644 index 000000000000..c320ac9e5ca0 --- /dev/null +++ b/hack/lint/backend_wrappers.go @@ -0,0 +1,46 @@ +//go:build ruleguard + +// Package gorules holds the go-ruleguard rules gocritic runs inside +// `make lint`. It is never compiled into the binary: the build tag keeps it out +// of every normal build, and golangci-lint loads the file as data. +package gorules + +import "github.com/quasilyte/go-ruleguard/dsl" + +// backendWrapperMustBeUnwrappable fires on a struct that decorates a gRPC +// backend by embedding the raw interface. +// +// This exists because the same defect shipped twice. A wrapper that embeds +// grpc.Backend inherits exactly the methods Backend declares and nothing else. +// grpc.DialErrorReporter is deliberately NOT on Backend, so a wrapped client +// silently stops answering "did the transport fail, or did the backend die", +// and every guard built on that answer reads nil. The consequence is not +// subtle: core/services/nodes and pkg/model delete replica rows and stop +// backends on that answer, so a wrapper that swallows it turns a momentary loss +// of route into fleet-wide model reclamation. +// +// The rule is SYNTACTIC, and deliberately so. The obvious formulation, "embeds +// a backend and has no Unwrap", cannot be written: HasMethod rejects inline +// signatures outright ("inline func signatures are not supported yet"), its +// method-reference form needs a package ruleguard's own typechecker can import +// and it cannot import this module, and Implements tests the VALUE method set +// while every Unwrap here would be on a pointer receiver. So instead of +// checking for the method, this checks for the shape that CANNOT lack it: +// grpc.WrappedBackend provides the same pass-through method set plus Unwrap, +// with a value receiver, so anything embedding it is transparent by +// construction. Forgetting is then not expressible rather than merely +// discouraged, which is the same move loopbackService makes in the worker. +// +// Test doubles are excluded by path in .golangci.yml: they embed a NIL backend +// to inherit the interface's method set, decorate nothing, and have no +// transport answer to forward. +func backendWrapperMustBeUnwrappable(m dsl.Matcher) { + m.Import("github.com/mudler/LocalAI/pkg/grpc") + + m.Match( + `type $w struct { $*_; grpc.Backend; $*_ }`, + `type $w struct { $*_; grpc.ControlBackend; $*_ }`, + `type $w struct { $*_; grpc.InferenceBackend; $*_ }`, + ). + Report(`$w decorates a gRPC backend by embedding the raw interface, so grpc.LastDialErrorOf cannot see through it and every transport-failure guard behind it reads nil, which deletes replica rows for workers that are merely unroutable. Embed grpc.WrappedBackend instead: it gives the same pass-through plus Unwrap. If $w decorates nothing, silence this with //nolint:gocritic and say so.`) +} diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 1126e39b7269..c37d2d677b59 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -88,6 +88,30 @@ type BackendUnwrapper interface { Unwrap() Backend } +// WrappedBackend is what a decorator embeds INSTEAD of a Backend. +// +// It provides the pass-through method set exactly as embedding the interface +// did, and it provides Unwrap, so a decorator built on it is transparent to +// LastDialErrorOf by CONSTRUCTION rather than by remembering. That is the whole +// design: the same defect shipped twice, both times because a wrapper inherited +// only what Backend declares and DialErrorReporter is deliberately not on +// Backend, so the transport answer every reaping guard depends on silently +// became nil. +// +// Forgetting is therefore no longer possible for anything that embeds this, and +// embedding the raw interface instead is caught by the ruleguard rule in +// hack/lint/. A compile-time assertion cannot do that job: it only fires for a +// type that already declares the intent, which is precisely the type that did +// not forget. +// +// Unwrap takes a VALUE receiver, which is safe because this holds one interface +// and no lock, and is what lets the value type of any embedder satisfy +// BackendUnwrapper. +type WrappedBackend struct{ Backend } + +// Unwrap exposes the decorated client. +func (w WrappedBackend) Unwrap() Backend { return w.Backend } + // maxBackendUnwrapDepth bounds the walk below. Three wrappers exist today and // they nest at most two deep; the bound is a guard against a cycle a future // wrapper could introduce, not a limit anything real approaches. diff --git a/pkg/model/connection_evicting_client.go b/pkg/model/connection_evicting_client.go index bde2333bfa61..2c9e262b8754 100644 --- a/pkg/model/connection_evicting_client.go +++ b/pkg/model/connection_evicting_client.go @@ -16,7 +16,7 @@ import ( // still returned to the caller — the NEXT request will trigger rescheduling // via SmartRouter. type ConnectionEvictingClient struct { - grpc.Backend + grpc.WrappedBackend modelID string evict func() once sync.Once @@ -24,15 +24,11 @@ type ConnectionEvictingClient struct { var _ grpc.BackendUnwrapper = (*ConnectionEvictingClient)(nil) -// Unwrap exposes the client this one decorates, so grpc.LastDialErrorOf can see -// past it. -func (c *ConnectionEvictingClient) Unwrap() grpc.Backend { return c.Backend } - func newConnectionEvictingClient(inner grpc.Backend, modelID string, evict func()) grpc.Backend { return &ConnectionEvictingClient{ - Backend: inner, - modelID: modelID, - evict: evict, + WrappedBackend: grpc.WrappedBackend{Backend: inner}, + modelID: modelID, + evict: evict, } } From ed9a4b6b52e4e46422f1ab6b2beda750332898c5 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 18:03:52 +0000 Subject: [PATCH 35/42] docs(grpc): withdraw the lint-cost claim, and make the one nolint tidy-proof I reported that enabling gocritic pushed make lint past 600s. That was wrong, and it was wrong in a way worth naming: those runs happened right after I changed pkg/grpc's and core/services/nodes' interfaces, so the Go build cache was cold for essentially the whole repository including every backend, and test suites were running concurrently on the same machine. I attributed a cold-cache full-repo typecheck under load to the linter I had just enabled, and raised it as a cost without ever timing it against a baseline. A number with no control is not a measurement. Measured properly, with the golangci cache cleaned before every run and isolated GOCACHE directories for the cold ones so the shared cache was not wiped: warm, base 15s then 7s and current 8s then 7s; cold, base 87s and current 78s running base first, base 136s and current 79s running current first. The spread between the two cold base runs is larger than any gap between base and current, so gocritic with only the ruleguard checker costs nothing measurable. So the rule stays, unscoped. Scoping it to core and pkg was the fallback for a cost that does not exist, and adding that configuration would buy nothing. The one override gets the protection it needs instead. InFlightTrackingClient's nolint is exactly the kind of thing a later reader tidies away, so it now opens by saying not to, and states what breaks rather than what is intended: WrappedBackend embeds the full Backend interface, so adopting it there would promote every inference method as untracked pass-through, the build would stay GREEN, and in-flight accounting would silently stop covering whatever was added next. WrappedBackend's own doc carries the counterpart warning so a reader arriving from either side finds it. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/nodes/inflight.go | 27 ++++++++++++++++++--------- pkg/grpc/backend.go | 6 ++++++ 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index 8bf98584a821..cc62782ceb42 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -27,16 +27,25 @@ import ( // interface therefore breaks this file's build (see the var assertion below) // until it is wrapped with track() - so a new inference path can't be added // without an in-flight accounting decision. -// The ruleguard rule wants grpc.WrappedBackend here, and this is the one -// decorator that cannot use it. Embedding ControlBackend rather than Backend is -// what produces the compile-time guarantee below: leave an InferenceBackend -// method unwrapped and this type stops satisfying grpc.Backend. WrappedBackend -// embeds the FULL interface, so adopting it would silently restore -// pass-through for every inference method and delete that guarantee. The -// transparency the rule protects is provided explicitly instead, by the Unwrap -// and the grpc.BackendUnwrapper assertion above. +// DO NOT "fix" this by embedding grpc.WrappedBackend, and do not delete the +// nolint below. Both look like tidy-ups and both silently remove a guarantee. // -//nolint:gocritic // embeds ControlBackend on purpose; Unwrap is declared explicitly +// The ruleguard rule in hack/lint/ asks every decorator to embed +// grpc.WrappedBackend, because that makes Unwrap structural. This is the one +// decorator that must not, and the reason is the paragraph above: embedding +// ControlBackend rather than Backend is exactly what forces every +// InferenceBackend method to be declared and tracked here, on pain of a build +// failure. grpc.WrappedBackend embeds the FULL Backend interface, so adopting +// it would promote every inference method as untracked pass-through, the build +// would stay green, and in-flight accounting would silently stop covering +// whatever was added next. +// +// The transparency the rule exists to protect is still provided, explicitly: +// the wrapped field, the Unwrap method and the grpc.BackendUnwrapper assertion +// above. A spec drives it (see the wrapper transport specs), so removing them +// reddens rather than merely regressing. +// +//nolint:gocritic // embeds ControlBackend deliberately; see the paragraph above before changing this type InFlightTrackingClient struct { grpc.ControlBackend // passthrough for control-plane / streaming-constructor methods inner grpc.InferenceBackend // tracked inference methods delegate here diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index c37d2d677b59..a0fdd2ad8815 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -107,6 +107,12 @@ type BackendUnwrapper interface { // Unwrap takes a VALUE receiver, which is safe because this holds one interface // and no lock, and is what lets the value type of any embedder satisfy // BackendUnwrapper. +// +// It is NOT for every decorator. Embedding this promotes the whole Backend +// surface as pass-through, so a decorator that deliberately embeds a NARROWER +// interface to force itself to handle each method (see +// nodes.InFlightTrackingClient) must keep doing that and declare Unwrap by +// hand; adopting this there would restore pass-through silently. type WrappedBackend struct{ Backend } // Unwrap exposes the decorated client. From 1cf847f29ef92f29c940558070c9908b00490b36 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 18:47:37 +0000 Subject: [PATCH 36/42] feat(distributed): stop workers listening, and stop them advertising A worker now opens no listener on a routable interface and states no endpoint at registration. Backend processes and the file-transfer server bind loopback, and the frontend reaches both through the tunnel the worker dials. The bind address is built from loopbackHost, the same constant the tunnel's grpc tag dials, so "the worker binds where its tunnel dials" is one fact in one place rather than two literals that can drift. All three advertisement sites are closed, not one: the registration body, RegisterNodeRequest, and the per-backend address in the install reply. That third one was hiding a live bug. stopModelExact refuses a stop whose ExpectedAddress does not match what the worker recorded for the process. The worker recorded 127.0.0.1:port; handleBackendInstall reported advertiseHost:port; the router stored the reported one and sent it straight back. On any worker whose advertise host was not 127.0.0.1, every acknowledged model stop failed with an address mismatch. Nothing caught it because the e2e harness set LOCALAI_ADVERTISE_ADDR=127.0.0.1, which made the rewrite a no-op. Removing the rewrite makes the two strings the same by construction. The brief was wrong about two of the four functions it called dead. effectiveBasePort is the base of the backend port allocator and resolveHTTPAddr is the file server's bind address; deleting them would have deleted the port allocator and the file server. Only the two advertise* helpers were dead, and addr_test.go is rewritten rather than deleted, because the port arithmetic it pinned still needs pinning. NodeModel.Address survives with a narrowed meaning and is renamed WorkerLocalAddress, along with the install reply field that feeds it. The frontend still has to say WHICH backend process on a worker it means, and the port in this string is how it says it: it travels as a stream target and the worker dials its own loopback. The gorm column and the json key stay "address", so neither a migration nor an API break rides along. Every fall-back to the node's address is gone. installBackendOnNode now errors when a worker reports success without naming one, because substituting the now-always-empty node address would name an empty target, and the worker refuses that as an invalid stream, which is classified as the worker answering about its backend. That is the "a present worker reads as something it is not" class this phase forbids. DistributedModelStore.Range had the same shape and was already wrong: it built each remote model's client from the node's base gRPC port, never the port a backend process listens on, so Free and Status went to the wrong place. It uses the replica's address now. BackendNode.Address and HTTPAddress are kept but made provably inert: no writer, no reader that acts on them, and Register force-clears both on re-registration so an upgraded worker's stale advertisement does not outlive its own upgrade in the API and the Nodes page. Dropping the columns is a ~90-site edit across the specs, the e2e suite, the MCP dto and the UI; it is recorded as a follow-up rather than folded in here. A persistent tunnel 401 still does not trigger re-registration, and now for a reason rather than a deferral. Register CLEARS the node's replica rows, so re-registering on a 401 would delete a live worker's rows on every retry, and under the name collision that causes the 401 the two workers would take turns doing it forever: a credential failure causing model reclamation. It also cannot fix the named cause, since a collision is indistinguishable from a restart. The 401 log now names both causes and says nothing can reach this worker, which is true only now that it has no listener. The container healthcheck did not break the way the brief expected, since the listener still exists on loopback and the probe runs inside the container. It did have a real #10987 defect that this change makes the common case: it read LOCALAI_SERVE_ADDR only, while effectiveBasePort reads LOCALAI_ADDR first, so a worker on a non-default base port was probed on 50050 and reported unhealthy while working. It follows the same precedence now. Docs, the compose file and the e2e harness are updated in step: no inbound rule or published port is needed for a worker, the two advertise variables are gone, the remaining address variables are read for their port only, the firewall-the-file-transfer-port warning is narrowed to the LOCALAI_HTTP_ADDR opt-out, and the upgrade-order note no longer claims the worker still listens. The Nodes page showed node.address, which is now always blank, so it shows the node id instead. Eight mutations, all red on a named spec, including reverting the loopback bind, re-adding the address to the registration body, restoring both node-address fall-backs, dropping the force-clear, storing the endpoint's address again, and un-fixing the healthcheck. One of them caught a defect in a spec I had just written: it asserted 200 where the endpoint returns 201, which went unnoticed because core/http/endpoints/localai is not on the task's verify list. It is run here. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/http/endpoints/localai/nodes.go | 27 ++--- core/http/endpoints/localai/nodes_test.go | 40 +++++-- .../src/components/nodes/NodePanel.jsx | 6 +- core/http/react-ui/src/pages/NodeDetail.jsx | 2 +- core/services/messaging/subjects.go | 15 ++- core/services/nodes/disk_headroom_test.go | 2 +- core/services/nodes/distributed_store.go | 16 ++- core/services/nodes/distributed_store_test.go | 51 +++++++-- core/services/nodes/health.go | 16 ++- core/services/nodes/health_mock_test.go | 12 ++ core/services/nodes/health_test.go | 14 +-- .../nodes/managers_distributed_test.go | 16 +-- core/services/nodes/model_router.go | 4 +- core/services/nodes/model_router_test.go | 14 ++- core/services/nodes/reconciler.go | 12 +- .../nodes/reconciler_busy_probe_test.go | 14 +-- .../nodes/reconciler_inflight_leak_test.go | 16 +-- core/services/nodes/reconciler_test.go | 36 +++--- .../nodes/reconciler_worker_processes_test.go | 14 +-- core/services/nodes/registry.go | 64 ++++++++--- core/services/nodes/registry_test.go | 36 +++++- .../nodes/revision_eligibility_test.go | 6 +- core/services/nodes/router.go | 58 +++++++--- .../nodes/router_eviction_alias_test.go | 2 +- .../nodes/router_eviction_selector_test.go | 2 +- .../services/nodes/router_load_budget_test.go | 2 +- core/services/nodes/router_load_job_test.go | 2 +- .../nodes/router_load_timeout_test.go | 2 +- core/services/nodes/router_reap_load_test.go | 2 +- .../services/nodes/router_reservation_test.go | 2 +- .../nodes/router_revision_lifecycle_test.go | 2 +- .../nodes/router_staging_context_test.go | 4 +- .../nodes/router_staging_deadline_test.go | 4 +- core/services/nodes/router_test.go | 70 +++++++++--- .../nodes/router_unreachable_worker_test.go | 4 +- core/services/nodes/unloader.go | 2 +- core/services/nodes/unloader_test.go | 6 +- core/services/worker/addr_test.go | 103 ++++++++++-------- core/services/worker/config.go | 22 ++-- core/services/worker/lifecycle.go | 26 ++--- core/services/worker/registration.go | 51 +++------ core/services/worker/supervisor.go | 39 +++++-- core/services/worker/tunnel.go | 9 +- core/services/worker/worker.go | 25 +++-- docker-compose.distributed.yaml | 21 ++-- docs/content/features/distributed-mode.md | 53 +++++---- scripts/build/healthcheck.sh | 23 +++- scripts/build/healthcheck_test.sh | 15 +++ tests/e2e/distributed/cluster/cluster.go | 6 +- .../distributed/model_config_revision_test.go | 2 +- 50 files changed, 635 insertions(+), 357 deletions(-) diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index f404555ac064..0571e7d7f8df 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -77,10 +77,14 @@ func GetNodeEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { // RegisterNodeRequest is the request body for registering a new worker node. type RegisterNodeRequest struct { - Name string `json:"name"` - NodeType string `json:"node_type,omitempty"` // "backend" (default) or "agent" - Address string `json:"address"` - HTTPAddress string `json:"http_address,omitempty"` + Name string `json:"name"` + NodeType string `json:"node_type,omitempty"` // "backend" (default) or "agent" + // No address and no http_address. A worker has no inbound endpoint to + // register: it holds one outbound tunnel to a frontend replica and every + // protocol the frontend speaks to it travels on that. An older worker still + // sends both keys and they are ignored, which is the intended outcome: + // storing them would put a dialable-looking endpoint back in the API for + // something nothing dials. Token string `json:"token,omitempty"` TotalVRAM uint64 `json:"total_vram,omitempty"` AvailableVRAM uint64 `json:"available_vram,omitempty"` @@ -142,22 +146,15 @@ 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 + // A backend worker no longer has to state an address; the tunnel it + // dials is what makes it reachable, and requiring one here would refuse + // exactly the workers this design is for. if req.Name == "" { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name is required")) } - if nodeType == nodes.NodeTypeBackend && req.Address == "" { - return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address is required for backend workers")) - } if len(req.Name) > 255 { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name exceeds 255 characters")) } - if len(req.Address) > 512 { - return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address exceeds 512 characters")) - } - if len(req.HTTPAddress) > 512 { - return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "http_address exceeds 512 characters")) - } // Hash the token for storage (if provided) var tokenHash string @@ -177,8 +174,6 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au node := &nodes.BackendNode{ Name: req.Name, NodeType: nodeType, - Address: req.Address, - HTTPAddress: req.HTTPAddress, TokenHash: tokenHash, TotalVRAM: req.TotalVRAM, AvailableVRAM: req.AvailableVRAM, diff --git a/core/http/endpoints/localai/nodes_test.go b/core/http/endpoints/localai/nodes_test.go index dababff38420..2255116aba8c 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -289,7 +289,11 @@ var _ = Describe("Node HTTP handlers", func() { Expect(errObj["message"]).To(ContainSubstring("exceeds 255 characters")) }) - It("returns 400 when address is missing for backend node type", func() { + It("registers a backend worker that states no address", func() { + // This used to be a 400. It is the shape every worker now + // registers with: it has no inbound endpoint, it holds one outbound + // tunnel, and refusing it here would refuse exactly the workers the + // tunnel exists for. e := echo.New() body := `{"name":"worker-no-addr"}` req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) @@ -299,13 +303,35 @@ var _ = Describe("Node HTTP handlers", func() { handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{}) Expect(handler(c)).To(Succeed()) - Expect(rec.Code).To(Equal(http.StatusBadRequest)) + Expect(rec.Code).To(Equal(http.StatusCreated)) - var resp map[string]any - Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) - errObj, ok := resp["error"].(map[string]any) - Expect(ok).To(BeTrue()) - Expect(errObj["message"]).To(ContainSubstring("address is required")) + stored, err := registry.GetByName(context.Background(), "worker-no-addr") + Expect(err).ToNot(HaveOccurred()) + Expect(stored.NodeType).To(Equal(nodes.NodeTypeBackend)) + Expect(stored.Address).To(BeEmpty()) + Expect(stored.HTTPAddress).To(BeEmpty()) + }) + + It("stores no address even when a worker still sends one", func() { + // An older worker keeps sending both keys. Storing them would put a + // dialable-looking endpoint back into the API and the Nodes page for + // something nothing dials, and would leave a reader of either one + // unsure which workers are reached how. + e := echo.New() + body := `{"name":"worker-legacy-addr","address":"10.0.0.9:50051","http_address":"10.0.0.9:50050"}` + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{}) + Expect(handler(c)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusCreated)) + + stored, err := registry.GetByName(context.Background(), "worker-legacy-addr") + Expect(err).ToNot(HaveOccurred()) + Expect(stored.Address).To(BeEmpty()) + Expect(stored.HTTPAddress).To(BeEmpty()) }) It("returns 400 when node_type is invalid", func() { diff --git a/core/http/react-ui/src/components/nodes/NodePanel.jsx b/core/http/react-ui/src/components/nodes/NodePanel.jsx index 623db00936bf..b73714ab9972 100644 --- a/core/http/react-ui/src/components/nodes/NodePanel.jsx +++ b/core/http/react-ui/src/components/nodes/NodePanel.jsx @@ -19,7 +19,11 @@ export default function NodePanel({ node, models = [], onApprove, onDrain, onRes
{node.name} - {node.address} + {/* A worker has no address to show: it holds an outbound tunnel and + binds nothing routable. Its id is what identifies it in routing + logs, so that is what an operator needs here. Pre-tunnel nodes + may still carry an address until they re-register. */} + {node.address || node.id}
e.stopPropagation()}> {node.status === 'pending' && ( diff --git a/core/http/react-ui/src/pages/NodeDetail.jsx b/core/http/react-ui/src/pages/NodeDetail.jsx index bff7db526ad2..52d3fea1d5f2 100644 --- a/core/http/react-ui/src/pages/NodeDetail.jsx +++ b/core/http/react-ui/src/pages/NodeDetail.jsx @@ -78,7 +78,7 @@ export default function NodeDetail() { navigate('/app/nodes')} className="link-plain">