fix(studio): stage workspace uploads outside the revision lock

With PostgreSQL and a filesystem blob store, Put and CommitUpper took
the global revision row lock before reading the upload. A slow or large
upload then blocked every other write and purge and kept a pool
connection for its whole duration.

Split the filesystem blob Put into Stage (read, hash, fsync a temp
file) and Publish (rename). The workspace stages first, takes the lock,
then publishes, so a purge still cannot remove a blob between
publication and the row that references it. Publish syncs only the
folder that gains the name; newly created folders are already synced by
mkdir.

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

core/services/studio/workspace/blob/blob.go

@@ -43,9 +43,12 @@ type Store interface {
 
 // Filesystem is an optional capability for authoritative direct filesystem blobs.
 // Object-store caches do not implement it. PostgreSQL uses it to coordinate
-// ingest/publication with purge automatically, including non-distributed owners.
+// publication with purge automatically, including non-distributed owners. It
+// requires Stager: the workspace must publish a blob inside its revision lock
+// and read the upload outside it.
 type Filesystem interface {
 	Store
+	Stager
 	IsFilesystem() bool
 }
 

core/services/studio/workspace/blob/fs.go

@@ -90,19 +90,35 @@ func (s *fsStore) mkdir(path string) error {
 	return s.syncDir(filepath.Dir(path))
 }
 func (s *fsStore) Put(ctx context.Context, r io.Reader, max int64) (string, int64, error) {
-	tempDir := filepath.Join(s.root, "tmp")
-	if err := s.mkdir(tempDir); err != nil {
-		return "", 0, err
-	}
-	f, err := os.CreateTemp(tempDir, "upload-*")
+	st, err := s.Stage(ctx, r, max)
 	if err != nil {
 		return "", 0, err
 	}
-	// Publication checks Sync and Close below; deferred cleanup is best-effort,
-	// including after a successful rename has removed the temporary name.
+	defer st.Discard()
+	if err := st.Publish(ctx); err != nil {
+		return "", 0, err
+	}
+	return st.SHA(), st.Size(), nil
+}
+
+// Stage reads r into a durable temporary file. The slow part of an upload
+// happens here; Publish only renames.
+func (s *fsStore) Stage(ctx context.Context, r io.Reader, max int64) (Staged, error) {
+	tempDir := filepath.Join(s.root, "tmp")
+	if err := s.mkdir(tempDir); err != nil {
+		return nil, err
+	}
+	f, err := os.CreateTemp(tempDir, "upload-*")
+	if err != nil {
+		return nil, err
+	}
+	st := &fsStaged{s: s, path: f.Name()}
+	// Close and Sync are checked below; this cleanup is best-effort.
 	defer func() {
 		_ = f.Close()
-		_ = os.Remove(f.Name())
+		if err != nil {
+			st.Discard()
+		}
 	}()
 	h := sha256.New()
 	src := io.Reader(contextReader{ctx: ctx, r: r})
@@ -110,57 +126,72 @@ func (s *fsStore) Put(ctx context.Context, r io.Reader, max int64) (string, int6
 		src = io.LimitReader(src, max+1)
 	}
 	if err = s.stage("write"); err != nil {
-		return "", 0, err
+		return nil, err
 	}
-	n, err := io.Copy(io.MultiWriter(f, h), src)
-	if err != nil {
-		return "", 0, err
+	if st.size, err = io.Copy(io.MultiWriter(f, h), src); err != nil {
+		return nil, err
 	}
-	if max > 0 && n > max {
-		return "", 0, workspace.ErrTooLarge
+	if max > 0 && st.size > max {
+		err = workspace.ErrTooLarge
+		return nil, err
 	}
 	if err = ctx.Err(); err != nil {
-		return "", 0, err
-	}
-	sha := hex.EncodeToString(h.Sum(nil))
-	final := shaPath(s.root, sha)
-	if err = s.mkdir(filepath.Dir(final)); err != nil {
-		return "", 0, err
-	}
-	if err = s.safe(final); err != nil {
-		return "", 0, err
+		return nil, err
 	}
+	st.sha = hex.EncodeToString(h.Sum(nil))
 	if err = f.Chmod(0o400); err != nil {
-		return "", 0, err
+		return nil, err
 	}
 	if err = s.stage("file sync"); err != nil {
-		return "", 0, err
+		return nil, err
 	}
 	if err = f.Sync(); err != nil {
-		return "", 0, err
+		return nil, err
 	}
 	if err = f.Close(); err != nil {
-		return "", 0, err
+		return nil, err
 	}
-	if err = s.stage("rename"); err != nil {
-		return "", 0, err
+	return st, nil
+}
+
+type fsStaged struct {
+	s    *fsStore
+	path string
+	sha  string
+	size int64
+}
+
+func (t *fsStaged) SHA() string { return t.sha }
+func (t *fsStaged) Size() int64 { return t.size }
+
+// Discard removes the temporary file. It is safe after Publish.
+func (t *fsStaged) Discard() { _ = os.Remove(t.path) }
+
+// Publish moves the staged file to its final name. The file is already
+// durable, so only the folder that gains the name needs a sync: a folder
+// created on the way is synced into its parent by mkdir, and a leftover
+// temporary name after a crash is harmless.
+func (t *fsStaged) Publish(ctx context.Context) error {
+	s := t.s
+	if err := ctx.Err(); err != nil {
+		return err
 	}
-	if err = os.Rename(f.Name(), final); err != nil {
-		return "", 0, err
+	final := shaPath(s.root, t.sha)
+	if err := s.mkdir(filepath.Dir(final)); err != nil {
+		return err
 	}
-	// Never remove final after a sync failure: another writer can own that name.
-	for p := filepath.Dir(final); ; p = filepath.Dir(p) {
-		if err = s.syncDir(p); err != nil {
-			return "", 0, err
-		}
-		if p == s.root {
-			break
-		}
+	if err := s.safe(final); err != nil {
+		return err
 	}
-	if err = s.syncDir(tempDir); err != nil {
-		return "", 0, err
+	if err := s.stage("rename"); err != nil {
+		return err
 	}
-	return sha, n, nil
+	// Rename replaces an existing blob with identical bytes. Never remove
+	// final after a sync failure: another writer can own that name.
+	if err := os.Rename(t.path, final); err != nil {
+		return err
+	}
+	return s.syncDir(filepath.Dir(final))
 }
 func (s *fsStore) Open(_ context.Context, sha string) (io.ReadSeekCloser, error) {
 	if !validSHA(sha) {

core/services/studio/workspace/blob/publication_internal_test.go

@@ -70,8 +70,9 @@ var _ = Describe("shared filesystem publication", func() {
 		}
 		_, _, err = s.Put(context.Background(), strings.NewReader("new shards"), 0)
 		Expect(err).NotTo(HaveOccurred())
-		// Creation of tmp/two shard parents, final/ancestor directories and tmp.
-		Expect(count).To(Equal(7))
+		// Creation of tmp and of the two shard folders, then the folder that
+		// gains the name. A finished file needs no sync of every ancestor.
+		Expect(count).To(Equal(4))
 	})
 
 	It("preserves concurrent publication after a failed sync", func() {
@@ -103,6 +104,38 @@ var _ = Describe("shared filesystem publication", func() {
 		Expect(string(got)).To(Equal(content))
 	})
 
+	It("keeps a staged blob invisible until it is published", func() {
+		st, err := NewFS(GinkgoT().TempDir())
+		Expect(err).NotTo(HaveOccurred())
+		ctx := context.Background()
+		staged, err := st.(Stager).Stage(ctx, strings.NewReader("staged"), 0)
+		Expect(err).NotTo(HaveOccurred())
+		ok, err := st.Exists(ctx, staged.SHA())
+		Expect(err).NotTo(HaveOccurred())
+		Expect(ok).To(BeFalse())
+		Expect(staged.Publish(ctx)).To(Succeed())
+		staged.Discard()
+		ok, err = st.Exists(ctx, staged.SHA())
+		Expect(err).NotTo(HaveOccurred())
+		Expect(ok).To(BeTrue())
+	})
+
+	It("drops a staged blob that is never published", func() {
+		root := GinkgoT().TempDir()
+		st, err := NewFS(root)
+		Expect(err).NotTo(HaveOccurred())
+		sg := NewStaging(st)
+		sha, _, err := sg.Put(context.Background(), strings.NewReader("abandoned"), 0)
+		Expect(err).NotTo(HaveOccurred())
+		sg.Close()
+		ok, err := st.Exists(context.Background(), sha)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(ok).To(BeFalse())
+		entries, err := os.ReadDir(filepath.Join(root, "tmp"))
+		Expect(err).NotTo(HaveOccurred())
+		Expect(entries).To(BeEmpty())
+	})
+
 	It("honors cancellation at the end of bounded ingest", func() {
 		s, err := NewFS(GinkgoT().TempDir())
 		Expect(err).NotTo(HaveOccurred())

core/services/studio/workspace/blob/stage.go

@@ -0,0 +1,76 @@
+package blob
+
+import (
+	"context"
+	"io"
+)
+
+// Staged is a blob whose bytes are read and hashed but not yet visible under
+// its name.
+type Staged interface {
+	SHA() string
+	Size() int64
+	// Publish makes the blob visible. It is quick: the bytes are already
+	// written. Calling it again after success is not supported.
+	Publish(ctx context.Context) error
+	// Discard drops the staged bytes. It is safe after Publish and safe to
+	// call more than once.
+	Discard()
+}
+
+// Stager is an optional capability of a Store that can split Put in two: the
+// slow read of the upload (Stage) and the quick publication (Staged.Publish).
+// The workspace stages outside its revision lock and publishes inside it.
+type Stager interface {
+	Stage(ctx context.Context, r io.Reader, max int64) (Staged, error)
+}
+
+// Staging is a Store whose Put only stages the bytes. Publish makes every
+// staged blob visible and Close drops the ones that were not published. A
+// store that is not a Stager is used as it is: its Put publishes at once and
+// Publish has nothing to do.
+type Staging struct {
+	Store
+	stager Stager
+	staged []Staged
+}
+
+// NewStaging wraps st.
+func NewStaging(st Store) *Staging {
+	s := &Staging{Store: st}
+	s.stager, _ = st.(Stager)
+	return s
+}
+
+// Put stages the bytes of r. The returned hash is not readable until Publish.
+func (s *Staging) Put(ctx context.Context, r io.Reader, max int64) (string, int64, error) {
+	if s.stager == nil {
+		return s.Store.Put(ctx, r, max)
+	}
+	st, err := s.stager.Stage(ctx, r, max)
+	if err != nil {
+		return "", 0, err
+	}
+	s.staged = append(s.staged, st)
+	return st.SHA(), st.Size(), nil
+}
+
+// Publish makes the staged blobs visible.
+func (s *Staging) Publish(ctx context.Context) error {
+	for _, st := range s.staged {
+		if err := st.Publish(ctx); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Close discards what Publish did not publish. Call it with defer.
+func (s *Staging) Close() {
+	for _, st := range s.staged {
+		st.Discard()
+	}
+	s.staged = nil
+}
+
+var _ Stager = (*fsStore)(nil)

core/services/studio/workspace/gormstore/files.go

@@ -8,6 +8,8 @@ import (
 	"io"
 	"strings"
 	"time"
+
+	"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
 )
 
 // precheckPut refuses a Put before any bytes are stored, when the actor has
@@ -45,24 +47,22 @@ func (s *Workspace) Put(ctx context.Context, a Actor, p string, r io.Reader, o P
 	if o.Title != "" {
 		title = o.Title
 	}
-	var w *writer
-	if s.sharedFilesystem {
-		w, err = s.beginWrite(ctx, a, title, ChangeUpload)
-		if err != nil {
-			return Node{}, err
-		}
-		defer w.rollback()
-	}
-	sum, size, err := s.blobs.Put(ctx, r, s.limit(o.MaxSize))
+	// The upload is read before the revision lock is taken, so a slow client
+	// cannot hold the lock that every writer and purge shares. The blob only
+	// becomes visible inside the lock, where a purge cannot race it.
+	st := blob.NewStaging(s.blobs)
+	defer st.Close()
+	sum, size, err := st.Put(ctx, r, s.limit(o.MaxSize))
 	if err != nil {
 		return Node{}, err
 	}
-	if w == nil {
-		w, err = s.beginWrite(ctx, a, title, ChangeUpload)
-		if err != nil {
-			return Node{}, err
-		}
-		defer w.rollback()
+	w, err := s.beginWrite(ctx, a, title, ChangeUpload)
+	if err != nil {
+		return Node{}, err
+	}
+	defer w.rollback()
+	if err := st.Publish(ctx); err != nil {
+		return Node{}, err
 	}
 	n, asked, err := w.put(segs, sum, size, o)
 	if err != nil {

core/services/studio/workspace/gormstore/lock_internal_test.go

@@ -0,0 +1,105 @@
+package gormstore
+
+import (
+	"context"
+	"io"
+	"path/filepath"
+	"strings"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/http/auth"
+	"github.com/mudler/LocalAI/core/services/studio/workspace"
+	"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
+)
+
+// slowBlobs is a filesystem blob store whose upload stays in progress until
+// release is closed. It signals started once a Stage or Put is under way.
+type slowBlobs struct {
+	blob.Store
+	started chan struct{}
+	release chan struct{}
+}
+
+func (b *slowBlobs) wait(ctx context.Context) error {
+	select {
+	case b.started <- struct{}{}:
+	default:
+	}
+	select {
+	case <-b.release:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (b *slowBlobs) Put(ctx context.Context, r io.Reader, max int64) (string, int64, error) {
+	if err := b.wait(ctx); err != nil {
+		return "", 0, err
+	}
+	return b.Store.Put(ctx, r, max)
+}
+
+func (b *slowBlobs) Stage(ctx context.Context, r io.Reader, max int64) (blob.Staged, error) {
+	if err := b.wait(ctx); err != nil {
+		return nil, err
+	}
+	return b.Store.(blob.Stager).Stage(ctx, r, max)
+}
+
+var _ = Describe("a slow upload on a shared filesystem", func() {
+	var (
+		s    *Workspace
+		slow *slowBlobs
+	)
+	// The admin may write anywhere; the lock is global, so who writes is moot.
+	alice := workspace.Actor{Principal: "admin", ChatID: "1"}
+	bob := workspace.Actor{Principal: "admin", ChatID: "2"}
+
+	BeforeEach(func() {
+		db, err := auth.OpenSQLiteFile(filepath.Join(GinkgoT().TempDir(), "ws.db"))
+		Expect(err).ToNot(HaveOccurred())
+		DeferCleanup(func() {
+			if raw, err := db.DB(); err == nil {
+				_ = raw.Close()
+			}
+		})
+		fs, err := blob.NewFS(filepath.Join(GinkgoT().TempDir(), "blobs"))
+		Expect(err).ToNot(HaveOccurred())
+		slow = &slowBlobs{Store: fs, started: make(chan struct{}, 1), release: make(chan struct{})}
+		s, err = New(db, slow, Options{Admin: "admin"})
+		Expect(err).ToNot(HaveOccurred())
+		// SQLite stands in for PostgreSQL here; the lock order is the same.
+		s.sharedFilesystem = true
+		// Never leave the upload blocked, even when a spec fails.
+		DeferCleanup(func() {
+			select {
+			case <-slow.release:
+			default:
+				close(slow.release)
+			}
+		})
+	})
+
+	It("does not hold the revision lock while Put reads the upload", func() {
+		done := make(chan error, 1)
+		go func() {
+			_, err := s.Put(context.Background(), alice, "/big.bin", strings.NewReader("slow"), workspace.PutOptions{})
+			done <- err
+		}()
+		Eventually(slow.started).Should(Receive())
+
+		other := make(chan error, 1)
+		go func() {
+			_, err := s.Mkdir(context.Background(), bob, "/mine")
+			other <- err
+		}()
+		Eventually(other, 3*time.Second).Should(Receive(BeNil()), "another user's write waits for a slow upload")
+
+		close(slow.release)
+		Eventually(done, 5*time.Second).Should(Receive(BeNil()))
+	})
+})

core/services/studio/workspace/gormstore/shared_fs_postgres_test.go

@@ -39,6 +39,35 @@ func (b *gateBlob) IsFilesystem() bool {
 	return ok && fs.IsFilesystem()
 }
 
+// Stage gates publication, the point where the blob is visible and the
+// revision lock is held but the row that refers to it is not written yet.
+func (b *gateBlob) Stage(ctx context.Context, r io.Reader, max int64) (blob.Staged, error) {
+	st, err := b.Store.(blob.Stager).Stage(ctx, r, max)
+	if err != nil {
+		return nil, err
+	}
+	return &gateStaged{Staged: st, b: b}, nil
+}
+
+type gateStaged struct {
+	blob.Staged
+	b *gateBlob
+}
+
+func (g *gateStaged) Publish(ctx context.Context) error {
+	if err := g.Staged.Publish(ctx); err != nil {
+		return err
+	}
+	g.b.mu.Lock()
+	entered, resume := g.b.uploaded, g.b.resume
+	g.b.mu.Unlock()
+	if entered != nil {
+		close(entered)
+		<-resume
+	}
+	return nil
+}
+
 func (b *gateBlob) Put(ctx context.Context, r io.Reader, max int64) (string, int64, error) {
 	sha, n, err := b.Store.Put(ctx, r, max)
 	b.mu.Lock()
@@ -275,26 +304,24 @@ var _ = Describe("PostgreSQL shared filesystem", func() {
 		}
 	})
 
-	It("cancels slow shared ingest and releases the revision lock", func() {
-		a, b, ba, _ := sharedPG()
+	It("leaves the revision lock free while a slow upload is read", func() {
+		a, b, _, _ := sharedPG()
 		request, cancel := context.WithCancel(ctx)
 		defer cancel()
 		reader := &cancelUpload{ctx: request, entered: make(chan struct{})}
 		done := make(chan error, 1)
 		go func() { _, err := a.Put(request, actor, "/slow", reader, workspace.PutOptions{}); done <- err }()
 		<-reader.entered
-		waiter := make(chan error, 1)
-		go func() { _, err := b.Mkdir(ctx, actor, "/after"); waiter <- err }()
-		waitSharedLock(ba.pool)
+		_, err := b.Mkdir(ctx, actor, "/after")
+		Expect(err).NotTo(HaveOccurred(), "a slow upload must not block other writers")
 		cancel()
 		Expect(<-done).To(MatchError(context.Canceled))
-		Expect(<-waiter).To(Succeed())
-		_, err := b.Stat(ctx, actor, "/slow")
+		_, err = b.Stat(ctx, actor, "/slow")
 		Expect(err).To(MatchError(workspace.ErrNotFound))
 	})
 
 	for _, kind := range []string{"Put", "CommitUpper"} {
-		It("automatically coordinates default PostgreSQL direct-FS ingest for "+kind, func() {
+		It("holds the revision lock from publication to metadata for "+kind, func() {
 			a, b, ba, _ := sharedPG()
 			ba.uploaded = make(chan struct{})
 			ba.resume = make(chan struct{})

core/services/studio/workspace/gormstore/upper.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"slices"
 
+	"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
 	"github.com/mudler/LocalAI/core/services/studio/workspace/internal/upperset"
 )
 
@@ -67,24 +68,21 @@ func (s *Workspace) CommitUpper(ctx context.Context, v View, a Actor, base int64
 	if err != nil {
 		return ChangeSet{}, nil, err
 	}
-	var w *writer
-	if s.sharedFilesystem {
-		w, err = s.beginWrite(ctx, a, o.Title, ChangeAgent)
-		if err != nil {
-			return ChangeSet{}, nil, err
-		}
-		defer w.rollback()
-	}
-	u, err := upperset.Read(ctx, v, pts, gone0, src, s.blobs, s.maxSize)
+	// Files are read before the revision lock is taken and published inside
+	// it; see Put.
+	st := blob.NewStaging(s.blobs)
+	defer st.Close()
+	u, err := upperset.Read(ctx, v, pts, gone0, src, st, s.maxSize)
 	if err != nil {
 		return ChangeSet{}, nil, err
 	}
-	if w == nil {
-		w, err = s.beginWrite(ctx, a, o.Title, ChangeAgent)
-		if err != nil {
-			return ChangeSet{}, nil, err
-		}
-		defer w.rollback()
+	w, err := s.beginWrite(ctx, a, o.Title, ChangeAgent)
+	if err != nil {
+		return ChangeSet{}, nil, err
+	}
+	defer w.rollback()
+	if err := st.Publish(ctx); err != nil {
+		return ChangeSet{}, nil, err
 	}
 	own, err := ownRevisions(ctx, w.tx, a, o.Own)
 	if err != nil {