diff --git a/core/application/application.go b/core/application/application.go index cf00e3952..4c081f9f5 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -182,7 +182,7 @@ type Application struct { studioTempDir string // removed at shutdown; set only in memory-only mode studioHostExec *hostexec.Runner // the chats' built-in bash on this machine studioFiles *files.Service - studioFilesIndex io.Closer // the JSON index of the workspace; nil when it uses the database + 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 authDB *gorm.DB diff --git a/core/application/studio_files.go b/core/application/studio_files.go index 2bf520dda..412c7e2ed 100644 --- a/core/application/studio_files.go +++ b/core/application/studio_files.go @@ -219,12 +219,16 @@ func addonSkillSync(f *files.Service) addons.SkillSync { // openWorkspaceIndex opens the index of the workspace: the GORM store on the // auth database when there is one, and the JSON store in a folder under root -// otherwise. The returned closer is nil for the shared database, which the -// application closes itself. +// otherwise. The closer ends the removals that wait for their grace and, for +// the JSON store, releases its folder. The shared database is closed by the +// application itself. func (a *Application) openWorkspaceIndex(root string, blobs blob.Store) (workspace.Store, io.Closer, error) { if a.authDB != nil { ws, err := gormstore.New(a.authDB, blobs, gormstore.Options{Admin: files.AdminPrincipal, MaxFileSize: studioMaxFileSize}) - return ws, nil, err + if err != nil { + return nil, nil, err + } + return ws, ws, nil } ws, err := jsonstore.New(filepath.Join(root, "workspace"), blobs, jsonstore.Options{Admin: files.AdminPrincipal, MaxFileSize: studioMaxFileSize}) if err != nil { diff --git a/core/services/studio/workspace/blob/blob.go b/core/services/studio/workspace/blob/blob.go index 18cba90c2..a48d677f6 100644 --- a/core/services/studio/workspace/blob/blob.go +++ b/core/services/studio/workspace/blob/blob.go @@ -181,39 +181,82 @@ func Grace(d time.Duration) time.Duration { return d } -// RemoveLater removes the blobs in shas from st after grace, but only those -// that inUse still reports as free at that moment. The wait covers a Put that -// found the blob already stored and has not yet written the row that refers to -// it. A grace of zero or less removes at once. A store without Remove keeps -// its blobs. -func RemoveLater(st Store, shas []string, grace time.Duration, inUse func(ctx context.Context, shas []string) (map[string]bool, error)) int { - rm, ok := st.(Remover) +// SweepTimeout bounds one removal pass: the wait for a connection and for a +// lock, the check and the removals together. +const SweepTimeout = 30 * time.Second + +// InUse tells which of shas something still refers to. A store whose answer +// must stay true until the blobs are gone, for example because the check +// holds a lock, returns a release func; RemoveLater calls it after the last +// removal. The release func may be nil. +type InUse func(ctx context.Context, shas []string) (used map[string]bool, release func(), err error) + +// Sweep removes unreferenced blobs from a Store after a purge. +type Sweep struct { + Store Store + // Grace is the wait before the removal. It covers a Put that found the + // blob already stored and has not yet written the row that refers to it. + // Zero or less removes at once. + Grace time.Duration + InUse InUse + // Life ends the removals that still wait for their grace, for example at + // shutdown. Nil means they always run. + Life context.Context +} + +// RemoveLater removes the blobs in shas, but only those that InUse reports as +// free at that moment. With no grace it runs now, in ctx, and returns how many +// blobs it removed. Otherwise it runs after the grace, outside ctx and bounded +// by Life, and returns the number of blobs it will check. A Store without +// Remove keeps its blobs. +func (sw Sweep) RemoveLater(ctx context.Context, shas []string) int { + rm, ok := sw.Store.(Remover) if !ok || len(shas) == 0 { return 0 } - run := func() int { - ctx := context.Background() - used, err := inUse(ctx, shas) - if err != nil { - xlog.Warn("workspace: could not check the blobs of a purge", "error", err) - return 0 - } - n := 0 - for _, sha := range shas { - if used[sha] { - continue - } - if err := rm.Remove(ctx, sha); err != nil { - xlog.Warn("workspace: could not remove a blob", "sha", sha, "error", err) - continue - } - n++ - } - return n + if sw.Grace <= 0 { + return sw.run(ctx, rm, shas) } - if grace <= 0 { - return run() + life := sw.Life + if life == nil { + life = context.Background() } - time.AfterFunc(grace, func() { run() }) + timer := time.NewTimer(sw.Grace) + go func() { + select { + case <-timer.C: + sw.run(life, rm, shas) + case <-life.Done(): + timer.Stop() + } + }() return len(shas) } + +func (sw Sweep) run(ctx context.Context, rm Remover, shas []string) int { + ctx, cancel := context.WithTimeout(ctx, SweepTimeout) + defer cancel() + used, release, err := sw.InUse(ctx, shas) + if release != nil { + defer release() + } + if err != nil { + xlog.Warn("workspace: could not check the blobs of a purge", "error", err) + return 0 + } + n := 0 + for _, sha := range shas { + if used[sha] { + continue + } + if ctx.Err() != nil { + break + } + if err := rm.Remove(ctx, sha); err != nil { + xlog.Warn("workspace: could not remove a blob", "sha", sha, "error", err) + continue + } + n++ + } + return n +} diff --git a/core/services/studio/workspace/blob/grace_test.go b/core/services/studio/workspace/blob/grace_test.go index b29433a89..90518f906 100644 --- a/core/services/studio/workspace/blob/grace_test.go +++ b/core/services/studio/workspace/blob/grace_test.go @@ -1,6 +1,9 @@ package blob_test import ( + "context" + "errors" + "sync" "time" . "github.com/onsi/ginkgo/v2" @@ -19,3 +22,66 @@ var _ = Describe("Grace", func() { Entry("positive is kept", 5*time.Second, 5*time.Second), ) }) + +// removerStore is a Store that can remove, counting what it removed. +type removerStore struct { + blob.Store + mu sync.Mutex + removed []string +} + +func (r *removerStore) Remove(_ context.Context, sha string) error { + r.mu.Lock() + defer r.mu.Unlock() + r.removed = append(r.removed, sha) + return nil +} + +func (r *removerStore) got() []string { + r.mu.Lock() + defer r.mu.Unlock() + return append([]string(nil), r.removed...) +} + +var _ = Describe("Sweep", func() { + var rs *removerStore + BeforeEach(func() { rs = &removerStore{} }) + + It("removes only what InUse reports as free and then releases", func() { + released := false + sw := blob.Sweep{Store: rs, InUse: func(context.Context, []string) (map[string]bool, func(), error) { + return map[string]bool{"a": true}, func() { released = true }, nil + }} + Expect(sw.RemoveLater(context.Background(), []string{"a", "b"})).To(Equal(1)) + Expect(rs.got()).To(Equal([]string{"b"})) + Expect(released).To(BeTrue()) + }) + + It("keeps the blobs when the check fails", func() { + sw := blob.Sweep{Store: rs, InUse: func(context.Context, []string) (map[string]bool, func(), error) { + return nil, nil, errors.New("down") + }} + Expect(sw.RemoveLater(context.Background(), []string{"a"})).To(BeZero()) + Expect(rs.got()).To(BeEmpty()) + }) + + It("removes after the grace, even when the caller is gone", func() { + caller, cancel := context.WithCancel(context.Background()) + cancel() + sw := blob.Sweep{Store: rs, Grace: 10 * time.Millisecond, InUse: func(context.Context, []string) (map[string]bool, func(), error) { + return nil, nil, nil + }} + sw.RemoveLater(caller, []string{"a"}) + Eventually(rs.got).Should(Equal([]string{"a"})) + }) + + It("drops a removal that still waits when Life ends", func() { + life, end := context.WithCancel(context.Background()) + sw := blob.Sweep{Store: rs, Grace: 100 * time.Millisecond, Life: life, InUse: func(context.Context, []string) (map[string]bool, func(), error) { + return nil, nil, nil + }} + sw.RemoveLater(context.Background(), []string{"a"}) + end() + Consistently(rs.got, 300*time.Millisecond).Should(BeEmpty()) + }) +}) diff --git a/core/services/studio/workspace/gormstore/purge.go b/core/services/studio/workspace/gormstore/purge.go index f988fd77b..eab1a80e0 100644 --- a/core/services/studio/workspace/gormstore/purge.go +++ b/core/services/studio/workspace/gormstore/purge.go @@ -6,10 +6,8 @@ import ( "context" "database/sql" "errors" - "github.com/mudler/xlog" "slices" "strings" - "time" "github.com/mudler/LocalAI/core/services/studio/workspace/blob" ) @@ -251,17 +249,27 @@ func (s *Workspace) purge(ctx context.Context, a Actor, dryRun bool, pick func(q for sha := range freed { shas = append(shas, sha) } - if s.sharedFilesystem { - s.removePGBlobs(ctx, shas) - } else { - blob.RemoveLater(s.blobs, shas, blob.Grace(s.grace), s.blobsInUse) - } + s.removeBlobs(ctx, shas) return res, nil } +// removeBlobs removes the blobs that nothing refers to any more. On a shared +// filesystem the removal runs in the caller's scope when there is no grace, +// and under the revision lock; see blobsInUseLocked. +func (s *Workspace) removeBlobs(ctx context.Context, shas []string) { + sw := blob.Sweep{Store: s.blobs, Grace: blob.Grace(s.grace), InUse: s.blobsInUse, Life: s.life} + if s.sharedFilesystem { + sw.InUse = s.blobsInUseLocked + } else { + ctx = context.Background() + } + sw.RemoveLater(ctx, shas) +} + // blobsInUse tells which of shas a node or a journal row refers to. -func (s *Workspace) blobsInUse(ctx context.Context, shas []string) (map[string]bool, error) { - return blobsInUse(ctx, s.db, shas) +func (s *Workspace) blobsInUse(ctx context.Context, shas []string) (map[string]bool, func(), error) { + used, err := blobsInUse(ctx, s.db, shas) + return used, nil, err } func blobsInUse(ctx context.Context, q querier, shas []string) (map[string]bool, error) { used := make(map[string]bool, len(shas)) @@ -292,52 +300,26 @@ func scanStrings(rows *sql.Rows, into map[string]bool) error { return rows.Err() } -// The revision row covers ingest through reference publication and this entire -// check/delete interval across PostgreSQL connections. Grace alone cannot do so. -func (s *Workspace) removePGBlobs(ctx context.Context, shas []string) { - rm, ok := s.blobs.(blob.Remover) - if !ok || len(shas) == 0 { - return - } - if grace := blob.Grace(s.grace); grace > 0 { - // Delayed work outlives the request; its timeout starts after grace. - time.AfterFunc(grace, func() { s.removePGBlobsNow(context.Background(), shas, rm) }) - } else { - s.removePGBlobsNow(ctx, shas, rm) - } -} - -// removePGBlobsNow keeps the check/delete transaction in one cancellation scope. -func (s *Workspace) removePGBlobsNow(caller context.Context, shas []string, rm blob.Remover) { - // Bound pool acquisition, revision-lock waiting and reference checks together. - ctx, cancel := context.WithTimeout(caller, 30*time.Second) - defer cancel() +// blobsInUseLocked is blobsInUse under the revision row lock, which stays held +// until release. A writer publishes a blob and the row that refers to it under +// the same lock, so no writer can reference a blob between this check and its +// removal on any PostgreSQL connection. Grace alone cannot do that. +func (s *Workspace) blobsInUseLocked(ctx context.Context, shas []string) (map[string]bool, func(), error) { t, err := s.db.BeginTx(ctx) if err != nil { - xlog.Warn("workspace: could not begin blob purge", "error", err) - return + return nil, nil, err } - defer func() { _ = t.Rollback() }() - if _, err = t.ExecContext(ctx, `UPDATE studio_ws_meta SET value = value WHERE key = 'rev'`); err != nil { - xlog.Warn("workspace: could not lock blob purge", "error", err) - return + // Nothing is written: the transaction only holds the lock, so release + // rolls it back. + release := func() { _ = t.Rollback() } + if _, err := t.ExecContext(ctx, `UPDATE studio_ws_meta SET value = value WHERE key = 'rev'`); err != nil { + release() + return nil, nil, err } used, err := blobsInUse(ctx, t, shas) if err != nil { - xlog.Warn("workspace: could not check the blobs of a purge", "error", err) - return - } - for _, sha := range shas { - if ctx.Err() != nil { - return - } - if !used[sha] { - if err := rm.Remove(ctx, sha); err != nil { - xlog.Warn("workspace: could not remove a purged blob", "error", err) - } - } - } - if err := t.Commit(); err != nil { - xlog.Warn("workspace: could not finish blob purge", "error", err) + release() + return nil, nil, err } + return used, release, nil } diff --git a/core/services/studio/workspace/gormstore/shared_fs_postgres_test.go b/core/services/studio/workspace/gormstore/shared_fs_postgres_test.go index f8658c10b..e2db9bba7 100644 --- a/core/services/studio/workspace/gormstore/shared_fs_postgres_test.go +++ b/core/services/studio/workspace/gormstore/shared_fs_postgres_test.go @@ -145,7 +145,7 @@ var _ = Describe("PostgreSQL shared filesystem", func() { cleanup, cancel := context.WithTimeout(ctx, 500*time.Millisecond) defer cancel() done := make(chan struct{}) - go func() { b.removePGBlobsNow(cleanup, []string{sha}, bb); close(done) }() + go func() { b.removeBlobs(cleanup, []string{sha}); close(done) }() waitSharedLock(ba.pool) Eventually(done, 2*time.Second).Should(BeClosed()) Eventually(func() int { return bb.pool.Stats().InUse }, time.Second).Should(BeZero()) @@ -160,7 +160,7 @@ var _ = Describe("PostgreSQL shared filesystem", func() { Expect(err).NotTo(HaveOccurred()) caller, cancel := context.WithCancel(ctx) cancel() - b.removePGBlobs(caller, []string{sha}) + b.removeBlobs(caller, []string{sha}) exists, err := bb.Store.Exists(ctx, sha) Expect(err).NotTo(HaveOccurred()) Expect(exists).To(BeTrue()) @@ -174,7 +174,7 @@ var _ = Describe("PostgreSQL shared filesystem", func() { b.grace = 20 * time.Millisecond caller, cancel := context.WithCancel(ctx) cancel() - b.removePGBlobs(caller, []string{sha}) + b.removeBlobs(caller, []string{sha}) Eventually(func() bool { exists, err := bb.Store.Exists(ctx, sha) Expect(err).NotTo(HaveOccurred()) diff --git a/core/services/studio/workspace/gormstore/store.go b/core/services/studio/workspace/gormstore/store.go index 840880dae..ef79440ff 100644 --- a/core/services/studio/workspace/gormstore/store.go +++ b/core/services/studio/workspace/gormstore/store.go @@ -53,6 +53,8 @@ type Workspace struct { now func() time.Time grace time.Duration sharedFilesystem bool + life context.Context // ends the removals that wait for their grace + endLife context.CancelFunc } var ( @@ -89,7 +91,16 @@ func New(db *gorm.DB, blobs blob.Store, o Options) (*Workspace, error) { now = time.Now } fs, direct := blobs.(blob.Filesystem) - return &Workspace{db: c, blobs: blobs, admin: o.Admin, maxSize: o.MaxFileSize, now: now, grace: o.BlobGrace, sharedFilesystem: direct && fs.IsFilesystem() && d.name == "postgres"}, nil + life, endLife := context.WithCancel(context.Background()) + return &Workspace{db: c, blobs: blobs, admin: o.Admin, maxSize: o.MaxFileSize, now: now, grace: o.BlobGrace, sharedFilesystem: direct && fs.IsFilesystem() && d.name == "postgres", life: life, endLife: endLife}, nil +} + +// Close ends the blob removals that still wait for their grace. It does not +// close the database, which the caller owns. The blobs stay; a later purge +// finds them again. +func (s *Workspace) Close() error { + s.endLife() + return nil } // limit returns the smaller non-zero of the store limit and a per-call limit. diff --git a/core/services/studio/workspace/jsonstore/purge.go b/core/services/studio/workspace/jsonstore/purge.go index 9b4a827ed..63a392b7c 100644 --- a/core/services/studio/workspace/jsonstore/purge.go +++ b/core/services/studio/workspace/jsonstore/purge.go @@ -67,7 +67,8 @@ func (s *Store) purge(ctx context.Context, a workspace.Actor, dryRun bool, pick return res, err } // Outside the lock: the check of the blobs takes the read lock. - blob.RemoveLater(s.blobs, freed, blob.Grace(s.blobGrace), s.blobsInUse) + blob.Sweep{Store: s.blobs, Grace: blob.Grace(s.blobGrace), InUse: s.blobsInUse, Life: s.life}. + RemoveLater(context.Background(), freed) return res, nil } @@ -180,9 +181,9 @@ func deletedSubtree(d *data, root int64) (ids []int64, whole bool) { } // blobsInUse tells which of shas a node or a journal row refers to. -func (s *Store) blobsInUse(ctx context.Context, shas []string) (map[string]bool, error) { +func (s *Store) blobsInUse(ctx context.Context, shas []string) (map[string]bool, func(), error) { if err := ctx.Err(); err != nil { - return nil, err + return nil, nil, err } want := make(map[string]bool, len(shas)) for _, sha := range shas { @@ -202,5 +203,5 @@ func (s *Store) blobsInUse(ctx context.Context, shas []string) (map[string]bool, mark(r.Old.Blob) mark(r.New.Blob) } - return want, nil + return want, nil, nil } diff --git a/core/services/studio/workspace/jsonstore/store.go b/core/services/studio/workspace/jsonstore/store.go index 0325de5e6..bf972327c 100644 --- a/core/services/studio/workspace/jsonstore/store.go +++ b/core/services/studio/workspace/jsonstore/store.go @@ -82,6 +82,8 @@ type Store struct { now func() time.Time compactAt int64 blobGrace time.Duration + life context.Context // ends the removals that wait for their grace + endLife context.CancelFunc } var ( @@ -102,6 +104,7 @@ func New(dir string, blobs blob.Store, o Options) (*Store, error) { return nil, errors.New("workspace: a blob store is required") } s := &Store{d: newData(), blobs: blobs, admin: o.Admin, maxSize: o.MaxFileSize, now: o.Now, compactAt: o.CompactAt, blobGrace: o.BlobGrace} + s.life, s.endLife = context.WithCancel(context.Background()) if s.now == nil { s.now = time.Now } @@ -133,6 +136,7 @@ func New(dir string, blobs blob.Store, o Options) (*Store, error) { // Close releases the folder. The store must not be used afterwards. func (s *Store) Close() error { + s.endLife() s.mu.Lock() defer s.mu.Unlock() if s.dir == "" {