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 <mudler@localai.io>

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")

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)
 	}
 

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

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"))
+	})
+})

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

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)

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) {

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))
 	})
 

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() {

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)
 	})