refactor(studio): one blob sweep for every workspace store

The PostgreSQL store carried its own copy of the delayed blob removal.
Its timer was never stopped, it ran under context.Background(), and the
30 second bound was a bare number.

Replace blob.RemoveLater with blob.Sweep. The reference check may now
return a release func, so the PostgreSQL check can keep the revision
lock until the last removal. A Life context ends pending removals, and
SweepTimeout names the bound. Both stores cancel Life on Close, and the
application now closes the PostgreSQL store with the JSON one.

Assisted-by: Claude Code:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

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

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 {

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

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

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
 }

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

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.

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
 }

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 == "" {