fix(studio): time turn heartbeats with the database clock

A turn beat was written with the clock of the frontend that runs the
turn and compared with the clock of the frontend that reaps. A skew of
about the turn timeout made the reaper end live turns.

ReapTurns now takes an idle duration, and the GORM store reads the time
from the database for the claim, the beat and the cutoff. The JSON store
is single-process and keeps the local clock.

Assisted-by: Claude Code:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-11-03 08:38:53 +00:00
1 parent fd808e4be0
commit cc0ab51438
8 files changed
+69 -18

No files matched your search

+1 -1
View File
@@ -436,7 +436,7 @@ var _ = Describe("Loops of a chat", func() {
until("the run to start", func() bool { return turn(a.svc, alice, chat).Running })
// The process dies: its host stops, and the reaper frees the chat.
a.host.Close()
_, err := sh.chats.ReapTurns(ctx, time.Now().Add(time.Hour))
_, err := sh.chats.ReapTurns(ctx, -time.Hour)
Expect(err).NotTo(HaveOccurred())
b := sh.rig(say, studio.LoopOptions{})
+1 -1
View File
@@ -1125,7 +1125,7 @@ func (s *Service) heartbeat(ctx context.Context, p, conv, runID string) {
// startup with a zero duration, because every claim then belongs to a process
// that no longer exists.
func (s *Service) ReapStale(ctx context.Context, olderThan time.Duration) (int, error) {
reaped, err := s.d.Store.ReapTurns(ctx, time.Now().Add(-olderThan))
reaped, err := s.d.Store.ReapTurns(ctx, olderThan)
for _, r := range reaped {
k := key(r.Principal, r.ChatID)
// The frontend that ran the turn is gone; so are the waiters of its
+38 -4
View File
@@ -894,10 +894,35 @@ func (s *gormStore) ConfirmPending(ctx context.Context, principal, id string, id
})
}
// now returns the time of the database server. A turn's beat is written by the
// frontend that runs it and read by the one that reaps, so both use this one
// clock: a skew between the machines must not end a live turn.
func (s *gormStore) now(ctx context.Context) (time.Time, error) {
switch s.db.Dialector.Name() {
case "sqlite":
var v string
if err := s.db.WithContext(ctx).Raw("SELECT strftime('%Y-%m-%d %H:%M:%f', 'now')").Scan(&v).Error; err != nil {
return time.Time{}, err
}
return time.ParseInLocation("2006-01-02 15:04:05.000", v, time.UTC)
case "postgres":
var t time.Time
if err := s.db.WithContext(ctx).Raw("SELECT now()").Scan(&t).Error; err != nil {
return time.Time{}, err
}
return t.UTC(), nil
}
return time.Now().UTC(), nil
}
func (s *gormStore) ClaimTurn(ctx context.Context, principal, id, runID string) error {
now, err := s.now(ctx)
if err != nil {
return err
}
db := s.db.WithContext(ctx)
res := db.Model(&chatRow{}).Where("id = ? AND principal = ? AND (run_id = ? OR run_id IS NULL)", id, principal, "").
Updates(map[string]any{"run_id": runID, "run_since": time.Now().UTC(), "run_beat": time.Now().UTC()})
Updates(map[string]any{"run_id": runID, "run_since": now, "run_beat": now})
if res.Error != nil {
return res.Error
}
@@ -939,9 +964,13 @@ func (s *gormStore) Turn(ctx context.Context, principal, id string) (TurnInfo, e
}
func (s *gormStore) TouchTurn(ctx context.Context, principal, id, runID string) error {
now, err := s.now(ctx)
if err != nil {
return err
}
db := s.db.WithContext(ctx)
res := db.Model(&chatRow{}).Where("id = ? AND principal = ? AND run_id = ?", id, principal, runID).
Update("run_beat", time.Now().UTC())
Update("run_beat", now)
if res.Error != nil {
return res.Error
}
@@ -954,9 +983,14 @@ func (s *gormStore) TouchTurn(ctx context.Context, principal, id, runID string)
return ErrNotClaimed
}
func (s *gormStore) ReapTurns(ctx context.Context, olderThan time.Time) ([]ReapedTurn, error) {
func (s *gormStore) ReapTurns(ctx context.Context, idle time.Duration) ([]ReapedTurn, error) {
now, err := s.now(ctx)
if err != nil {
return nil, err
}
olderThan := now.Add(-idle)
var rows []chatRow
err := s.db.WithContext(ctx).
err = s.db.WithContext(ctx).
Where("run_id <> ? AND run_id IS NOT NULL AND run_beat < ?", "", olderThan).
Find(&rows).Error
if err != nil {
+18 -1
View File
@@ -60,6 +60,23 @@ var _ = Describe("the GORM store on SQLite", func() {
storetest.CheckSeed(open())
})
// The beat of a turn and the cutoff of the reaper come from the database
// clock, so the reap does not depend on the clock of the calling machine.
It("measures how long a turn is idle on the database clock", func() {
ctx := context.Background()
db := openSQLite(filepath.Join(GinkgoT().TempDir(), "studio.db"))
s, err := store.NewGORM(db)
Expect(err).NotTo(HaveOccurred())
now := time.Now().UTC()
Expect(s.CreateChat(ctx, session.ChatRecord{ID: "a", Principal: "alice", Created: now, Updated: now})).To(Succeed())
Expect(s.ClaimTurn(ctx, "alice", "a", "r1")).To(Succeed())
Expect(db.Exec("UPDATE studio_chats SET run_beat = ? WHERE id = ?", now.Add(-10*time.Minute), "a").Error).To(Succeed())
Expect(s.ReapTurns(ctx, 15*time.Minute)).To(BeEmpty(), "a turn idle for 10 minutes was reaped at 15")
Expect(s.TouchTurn(ctx, "alice", "a", "r1")).To(Succeed())
Expect(db.Exec("UPDATE studio_chats SET run_beat = ? WHERE id = ?", now.Add(-10*time.Minute), "a").Error).To(Succeed())
Expect(s.ReapTurns(ctx, 5*time.Minute)).To(HaveLen(1))
})
// ReapTurns frees the chat and restores its messages in one transaction: when
// the restore fails the chat must stay claimed, or the lost messages could
// never be given back.
@@ -75,7 +92,7 @@ var _ = Describe("the GORM store on SQLite", func() {
Expect(s.TakePending(ctx, "alice", "a")).Error().NotTo(HaveOccurred())
Expect(db.Exec(`CREATE TRIGGER no_restore BEFORE UPDATE ON studio_pending
BEGIN SELECT RAISE(ABORT, 'forced'); END`).Error).NotTo(HaveOccurred())
Expect(s.ReapTurns(ctx, time.Now().Add(time.Hour))).Error().To(HaveOccurred(), "ReapTurns should fail when the restore fails")
Expect(s.ReapTurns(ctx, -time.Hour)).Error().To(HaveOccurred(), "ReapTurns should fail when the restore fails")
info, err := s.Turn(ctx, "alice", "a")
Expect(err).NotTo(HaveOccurred())
Expect(info.Running).To(BeTrue(), "turn after failed reap = %+v; want still running", info)
+2 -1
View File
@@ -999,7 +999,8 @@ func (s *jsonStore) TouchTurn(_ context.Context, principal, id, runID string) er
return nil
}
func (s *jsonStore) ReapTurns(_ context.Context, olderThan time.Time) ([]ReapedTurn, error) {
func (s *jsonStore) ReapTurns(_ context.Context, idle time.Duration) ([]ReapedTurn, error) {
olderThan := time.Now().UTC().Add(-idle)
var out []ReapedTurn
for _, c := range s.slots() {
reaped, err := s.reapTurn(c, olderThan)
+4 -3
View File
@@ -168,11 +168,12 @@ type Store interface {
// TouchTurn records that the run holding the chat is still alive. The
// frontend that runs a turn calls it on a timer.
TouchTurn(ctx context.Context, principal, chatID, runID string) error
// ReapTurns ends every running turn whose last sign of life is before
// olderThan: the chat is free again, its last outcome is interrupted, and
// ReapTurns ends every running turn that has not beaten for idle, measured
// on the clock of the store, so that frontends with different clocks agree:
// the chat is free again, its last outcome is interrupted, and
// the messages a request took but nobody confirmed are undelivered again. A
// frontend that died leaves such turns behind. It returns the turns it ended.
ReapTurns(ctx context.Context, olderThan time.Time) ([]ReapedTurn, error)
ReapTurns(ctx context.Context, idle time.Duration) ([]ReapedTurn, error)
// Turn returns the turn state of the chat.
Turn(ctx context.Context, principal, chatID string) (TurnInfo, error)
}
@@ -580,12 +580,10 @@ func Run(newStore New) {
chat(s, "alive", "alice")
Expect(s.ClaimTurn(ctx(), "alice", "stale", "r-stale")).To(Succeed())
Expect(s.ClaimTurn(ctx(), "alice", "alive", "r-alive")).To(Succeed())
time.Sleep(40 * time.Millisecond)
cutoff := time.Now()
time.Sleep(40 * time.Millisecond)
time.Sleep(100 * time.Millisecond)
Expect(s.TouchTurn(ctx(), "alice", "alive", "r-alive")).To(Succeed())
got, err := s.ReapTurns(ctx(), cutoff)
got, err := s.ReapTurns(ctx(), 50*time.Millisecond)
Expect(err).NotTo(HaveOccurred())
Expect(got).To(HaveLen(1), "reaped = %+v, want only the chat that stopped beating", got)
Expect(got[0].ChatID).To(Equal("stale"))
@@ -599,7 +597,7 @@ func Run(newStore New) {
// The run that was reaped can no longer beat or release, and the chat is free.
Expect(s.TouchTurn(ctx(), "alice", "stale", "r-stale")).To(MatchError(store.ErrNotClaimed), "TouchTurn on a reaped run")
Expect(s.ClaimTurn(ctx(), "alice", "stale", "r-next")).To(Succeed(), "claim after reap")
again, _ := s.ReapTurns(ctx(), cutoff)
again, _ := s.ReapTurns(ctx(), 50*time.Millisecond)
Expect(again).To(BeEmpty(), "a second reap found %+v", again)
})
@@ -611,7 +609,7 @@ func Run(newStore New) {
kept, _ := s.AddPending(ctx(), "alice", "c1", "taken and recorded")
Expect(s.TakePending(ctx(), "alice", "c1")).Error().NotTo(HaveOccurred())
Expect(s.ConfirmPending(ctx(), "alice", "c1", []string{kept.ID})).To(Succeed())
got, _ := s.ReapTurns(ctx(), time.Now().Add(time.Hour))
got, _ := s.ReapTurns(ctx(), -time.Hour)
Expect(got).To(HaveLen(1), "reaped = %+v, want the one turn", got)
pend, _ := s.Pending(ctx(), "alice", "c1")
Expect(pend).To(HaveLen(1), "pending after the reap = %+v, want only the unconfirmed message, undelivered", pend)
+1 -1
View File
@@ -494,7 +494,7 @@ var _ = Describe("triggersvc", func() {
})
// The frontend that runs it dies: its host stops, and the reaper frees its chat.
Expect(a.sh.chats).NotTo(BeNil(), "no store")
reaped, err := sh.chats.ReapTurns(ctx, sh.clock.Now().Add(24*time.Hour).Add(time.Now().Sub(sh.clock.Now())))
reaped, err := sh.chats.ReapTurns(ctx, -time.Hour)
Expect(err).NotTo(HaveOccurred())
Expect(reaped).NotTo(BeEmpty(), "the reaper found no running turn")
// Whichever host supervised the run, the next pass of the new leader ends it.