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

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

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()
-		Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed())
-		reg = cluster.NewRegistry(db)
 		ctx = context.Background()
+		Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed())
+		Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed())
+		reg = cluster.NewRegistry(db)
 	})
 
 	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

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