diff --git a/core/services/studio/workspace/blob/blob.go b/core/services/studio/workspace/blob/blob.go index 0f1e1b6476..18cba90c28 100644 --- a/core/services/studio/workspace/blob/blob.go +++ b/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 } diff --git a/core/services/studio/workspace/blob/fs.go b/core/services/studio/workspace/blob/fs.go index 7a6a6df7d6..686c7eb501 100644 --- a/core/services/studio/workspace/blob/fs.go +++ b/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) { diff --git a/core/services/studio/workspace/blob/publication_internal_test.go b/core/services/studio/workspace/blob/publication_internal_test.go index 6300c3ca48..46805e5d33 100644 --- a/core/services/studio/workspace/blob/publication_internal_test.go +++ b/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()) diff --git a/core/services/studio/workspace/blob/stage.go b/core/services/studio/workspace/blob/stage.go new file mode 100644 index 0000000000..3b33c64d24 --- /dev/null +++ b/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) diff --git a/core/services/studio/workspace/gormstore/files.go b/core/services/studio/workspace/gormstore/files.go index e7653b2bfe..d76b7cb49f 100644 --- a/core/services/studio/workspace/gormstore/files.go +++ b/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 { diff --git a/core/services/studio/workspace/gormstore/lock_internal_test.go b/core/services/studio/workspace/gormstore/lock_internal_test.go new file mode 100644 index 0000000000..1d9e55a3b5 --- /dev/null +++ b/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())) + }) +}) 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 6ccd8c1f52..f8658c10b2 100644 --- a/core/services/studio/workspace/gormstore/shared_fs_postgres_test.go +++ b/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{}) diff --git a/core/services/studio/workspace/gormstore/upper.go b/core/services/studio/workspace/gormstore/upper.go index 44d7b3a498..303ca6964c 100644 --- a/core/services/studio/workspace/gormstore/upper.go +++ b/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 {