diff --git a/core/application/application.go b/core/application/application.go index 026ad279df..0e4bad5ecf 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -203,6 +203,7 @@ type Application struct { studioFilesIndex io.Closer // closes the workspace index: ends pending blob removals, frees the JSON folder studioFilesTemp string // blob cache folder of a cluster with no data path; removed at shutdown studioTextIndex *textindex.Indexer // the text of the documents in the files; nil when files are off + studioTextIndexStop func() // stops the indexer's loop and waits for it; nil when it does not run authDB *gorm.DB metricsService *monitoring.LocalAIMetricsService statsRecorder *billing.Recorder @@ -727,6 +728,13 @@ func (a *Application) Shutdown() error { if a.studioTempDir != "" { _ = os.RemoveAll(a.studioTempDir) } + // Before the folders go: a pass in flight would write into them. + if a.studioTextIndexStop != nil { + a.studioTextIndexStop() + } + if c, ok := a.studioStore.(io.Closer); ok { + _ = c.Close() // ends the background indexing of earlier chats + } if a.studioFilesIndex != nil { _ = a.studioFilesIndex.Close() } diff --git a/core/application/studio_files.go b/core/application/studio_files.go index 40241708ce..0ce1aa2664 100644 --- a/core/application/studio_files.go +++ b/core/application/studio_files.go @@ -160,7 +160,13 @@ func (a *Application) initStudioTextIndex(f *files.Service, root string) { if cfg.Distributed.ChatWorker { return } - go ix.Run(a.backgroundContext()) + ctx, cancel := context.WithCancel(a.backgroundContext()) + done := make(chan struct{}) + a.studioTextIndexStop = func() { cancel(); <-done } + go func() { + defer close(done) + ix.Run(ctx) + }() } // syncAddonSkills lets the add-on owner keep the add-on skill layer in step diff --git a/core/services/studio/changemode_test.go b/core/services/studio/changemode_test.go index def972db97..daf8ba7dbc 100644 --- a/core/services/studio/changemode_test.go +++ b/core/services/studio/changemode_test.go @@ -476,6 +476,9 @@ var _ = Describe("how a memory note is saved", func() { Expect(r.conns.Ingest(ctx, "telegram", contract.ConnectorInbound{BindingID: b.ID, ConnectorMessage: contract.ConnectorMessage{ Chat: "11", Direct: true, SenderID: "11", SenderName: "S", Text: "remember"}})).To(Succeed()) waitFor("the answer", func() bool { return len(bot.answers()) >= 1 }) + // The harness records a tool result after the turn's events went + // out, so the answer can reach the bot first: wait for the result. + waitFor("the tool result", func() bool { return len(r.h.ResultsFor("propose_change")) >= 1 }) res := r.h.ResultsFor("propose_change") Expect(res).To(HaveLen(1)) Expect(res[0].Err).To(BeNil(), "%+v", res) diff --git a/core/services/studio/store/gorm.go b/core/services/studio/store/gorm.go index 122055cc02..23bd14e42a 100644 --- a/core/services/studio/store/gorm.go +++ b/core/services/studio/store/gorm.go @@ -139,6 +139,22 @@ func (pendingRow) TableName() string { return "studio_pending" } type gormStore struct { db *gorm.DB text textIndex + + // The backfill of the chat search index runs in the background: Close + // stops it and waits, so nothing writes once the owner of the database + // (or of its folder) is gone. + stopBackfill context.CancelFunc + backfillDone chan struct{} +} + +// Close stops the background indexing of earlier chats and waits for it. It +// does not close the database, which the caller owns. +func (s *gormStore) Close() error { + if s.stopBackfill != nil { + s.stopBackfill() + <-s.backfillDone + } + return nil } // NewGORM returns a Store on db (SQLite or PostgreSQL) and creates its tables. @@ -158,7 +174,13 @@ func NewGORM(db *gorm.DB) (Store, error) { // Indexing is true from here, so a search never reports the index as // complete before the first look at the history. s.text.indexing.Store(true) - go s.backfillTexts() + var ctx context.Context + ctx, s.stopBackfill = context.WithCancel(context.Background()) + s.backfillDone = make(chan struct{}) + go func() { + defer close(s.backfillDone) + s.backfillTexts(ctx) + }() return s, nil } diff --git a/core/services/studio/store/gormsearch.go b/core/services/studio/store/gormsearch.go index 608e7a273e..ea07245559 100644 --- a/core/services/studio/store/gormsearch.go +++ b/core/services/studio/store/gormsearch.go @@ -338,9 +338,8 @@ func (s *gormStore) titleChats(ctx context.Context, q chatsearch.Query, terms [] // from the chat it reached. A message written while it runs is indexed by the // write itself, and a row that backfill adds for a message that was just // removed is ignored by every search (it joins the entry). -func (s *gormStore) backfillTexts() { +func (s *gormStore) backfillTexts(ctx context.Context) { defer s.text.indexing.Store(false) - ctx := context.Background() var st textStateRow if err := s.db.Where("id = ?", 1).Limit(1).Find(&st).Error; err != nil { xlog.Warn("Chat search: could not read the backfill state", "error", err) @@ -359,6 +358,9 @@ func (s *gormStore) backfillTexts() { } last, started, logged := st.LastChat, time.Now(), time.Now() for { + if ctx.Err() != nil { + return + } var ids []string if err := s.db.Model(&chatRow{}).Where("id > ?", last).Order("id").Limit(backfillStep).Pluck("id", &ids).Error; err != nil { xlog.Debug("Chat search: backfill stopped", "error", err) @@ -368,6 +370,9 @@ func (s *gormStore) backfillTexts() { break } for _, id := range ids { + if ctx.Err() != nil { + return + } if err := s.backfillChat(ctx, id); err != nil { xlog.Debug("Chat search: backfill stopped", "error", err) return @@ -383,7 +388,11 @@ func (s *gormStore) backfillTexts() { logged = time.Now() xlog.Info("Chat search: indexing earlier chats", "done", s.text.done.Load(), "of", total) } - time.Sleep(10 * time.Millisecond) // leave the database to the chats being written + select { // leave the database to the chats being written + case <-ctx.Done(): + return + case <-time.After(10 * time.Millisecond): + } } // Rows of a chat deleted while the backfill read it. _ = s.db.Exec(`DELETE FROM studio_entry_texts WHERE NOT EXISTS (SELECT 1 FROM studio_chats c WHERE c.id = studio_entry_texts.chat_id)`).Error