diff --git a/core/services/worker/ephemeral_capacity.go b/core/services/worker/ephemeral_capacity.go new file mode 100644 index 000000000..32ba77574 --- /dev/null +++ b/core/services/worker/ephemeral_capacity.go @@ -0,0 +1,628 @@ +// SPDX-License-Identifier: MIT + +package worker + +import ( + "fmt" + "io" + "math" + "os" + "path/filepath" + "strings" + "sync" + + "github.com/mudler/LocalAI/pkg/xsysinfo" +) + +const ephemeralCapacityWriteChunk int64 = 64 << 10 + +// EphemeralCapacityError reports the values used to reject a reservation. +type EphemeralCapacityError struct { + RequestedBytes int64 + UsageBytes int64 + LimitBytes int64 + AvailableBytes int64 + HeadroomBytes int64 +} + +func (e *EphemeralCapacityError) Error() string { + return fmt.Sprintf( + "ephemeral capacity exceeded: requested=%d usage=%d limit=%d available=%d headroom=%d", + e.RequestedBytes, + e.UsageBytes, + e.LimitBytes, + e.AvailableBytes, + e.HeadroomBytes, + ) +} + +// EphemeralReservationConflictError reports an attempt to change an active +// reservation without first committing or releasing it. +type EphemeralReservationConflictError struct { + Path string + ActiveBytes int64 + RequestedBytes int64 +} + +func (e *EphemeralReservationConflictError) Error() string { + return fmt.Sprintf( + "ephemeral path %q already has an active reservation of %d bytes; requested %d bytes", + e.Path, + e.ActiveBytes, + e.RequestedBytes, + ) +} + +type ephemeralCapacityState uint8 + +const ( + ephemeralCapacityExisting ephemeralCapacityState = iota + ephemeralCapacityActive + ephemeralCapacityWriting + ephemeralCapacityCommitted +) + +type ephemeralCapacityEntry struct { + state ephemeralCapacityState + baseline int64 + reserved int64 + pending int64 + inflight int64 + openWriters int +} + +// EphemeralCapacityGuard accounts files and in-flight writes below a fixed set +// of ephemeral roots. +type EphemeralCapacityGuard struct { + mu sync.Mutex + changed *sync.Cond + roots []string + byteLimit int64 + minFreeBytes int64 + usage int64 + entries map[string]ephemeralCapacityEntry + commitWaiters map[string]int +} + +// NewEphemeralCapacityGuard creates a guard and accounts regular files already +// present below roots. Directory walks do not follow symbolic links. +func NewEphemeralCapacityGuard(roots []string, byteLimit, minFreeBytes int64) (*EphemeralCapacityGuard, error) { + if len(roots) == 0 { + return nil, fmt.Errorf("at least one ephemeral root is required") + } + if byteLimit < 0 { + return nil, fmt.Errorf("ephemeral byte limit must not be negative") + } + if minFreeBytes < 0 { + return nil, fmt.Errorf("ephemeral free-space headroom must not be negative") + } + + guard := &EphemeralCapacityGuard{ + roots: make([]string, 0, len(roots)), + byteLimit: byteLimit, + minFreeBytes: minFreeBytes, + entries: make(map[string]ephemeralCapacityEntry), + commitWaiters: make(map[string]int), + } + guard.changed = sync.NewCond(&guard.mu) + seenRoots := make(map[string]struct{}, len(roots)) + for _, root := range roots { + cleanRoot, err := cleanEphemeralAbsolutePath(root) + if err != nil { + return nil, fmt.Errorf("resolving ephemeral root: %w", err) + } + if _, found := seenRoots[cleanRoot]; found { + continue + } + if err := rejectEphemeralSymlinkComponents(cleanRoot); err != nil { + return nil, fmt.Errorf("validating ephemeral root %q: %w", cleanRoot, err) + } + seenRoots[cleanRoot] = struct{}{} + guard.roots = append(guard.roots, cleanRoot) + } + + for _, root := range guard.roots { + if err := guard.accountExistingFiles(root); err != nil { + return nil, err + } + } + return guard, nil +} + +// Reserve atomically reserves size additional bytes for path. An existing or +// committed file remains charged until the new write is committed. +func (g *EphemeralCapacityGuard) Reserve(path string, size int64) error { + if size < 0 { + return fmt.Errorf("reservation size must not be negative") + } + cleanPath, root, err := g.registeredPath(path) + if err != nil { + return err + } + + g.mu.Lock() + defer g.mu.Unlock() + entry, found := g.entries[cleanPath] + if found && entry.isActive() { + if entry.reserved == size { + return nil + } + return &EphemeralReservationConflictError{ + Path: cleanPath, ActiveBytes: entry.reserved, RequestedBytes: size, + } + } + if err := g.checkCapacityLocked(root, size); err != nil { + return err + } + entry.state = ephemeralCapacityActive + entry.reserved = size + entry.pending = size + entry.inflight = 0 + g.entries[cleanPath] = entry + g.usage += size + return nil +} + +// Commit replaces the path's baseline and reservation with the regular file's +// actual size. It waits for every bounded writer for the path to close. +func (g *EphemeralCapacityGuard) Commit(path string) error { + cleanPath, root, err := g.registeredPath(path) + if err != nil { + return err + } + + g.mu.Lock() + defer g.mu.Unlock() + for { + entry, found := g.entries[cleanPath] + if !found { + return fmt.Errorf("ephemeral path %q has no reservation", cleanPath) + } + if entry.state == ephemeralCapacityExisting || entry.state == ephemeralCapacityCommitted { + return nil + } + if entry.openWriters == 0 { + break + } + g.commitWaiters[cleanPath]++ + g.changed.Wait() + g.commitWaiters[cleanPath]-- + if g.commitWaiters[cleanPath] == 0 { + delete(g.commitWaiters, cleanPath) + } + } + + info, err := os.Lstat(cleanPath) + if err != nil { + return fmt.Errorf("stating committed ephemeral file %q: %w", cleanPath, err) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("committed ephemeral path %q is not a regular file", cleanPath) + } + entry := g.entries[cleanPath] + charged := entry.baseline + entry.reserved + if info.Size() > charged { + additional := info.Size() - charged + if err := g.checkCapacityLocked(root, additional); err != nil { + return err + } + charged += additional + g.usage += additional + } + g.usage -= charged - info.Size() + entry.state = ephemeralCapacityCommitted + entry.baseline = info.Size() + entry.reserved = 0 + entry.pending = 0 + entry.inflight = 0 + g.entries[cleanPath] = entry + return nil +} + +// Release forgets all accounting for path. It is safe to call repeatedly. +func (g *EphemeralCapacityGuard) Release(path string) error { + cleanPath, _, err := g.registeredPath(path) + if err != nil { + return err + } + + g.mu.Lock() + defer g.mu.Unlock() + for { + entry, found := g.entries[cleanPath] + if !found { + return nil + } + if entry.openWriters == 0 { + g.releaseLocked(cleanPath) + return nil + } + g.changed.Wait() + } +} + +// Account records bytes found by recovery after startup. Existing active +// reservations are left unchanged so recovery cannot erase in-flight charges. +func (g *EphemeralCapacityGuard) Account(path string, size int64) error { + if size < 0 { + return fmt.Errorf("accounted size must not be negative") + } + cleanPath, _, err := g.registeredPath(path) + if err != nil { + return err + } + + g.mu.Lock() + defer g.mu.Unlock() + entry, found := g.entries[cleanPath] + if found && entry.isActive() { + return &EphemeralReservationConflictError{ + Path: cleanPath, ActiveBytes: entry.reserved, RequestedBytes: size, + } + } + newUsage := g.usage + if found { + newUsage -= entry.baseline + } + if size > math.MaxInt64-newUsage { + return fmt.Errorf("ephemeral usage exceeds supported size") + } + entry = ephemeralCapacityEntry{state: ephemeralCapacityExisting, baseline: size} + g.entries[cleanPath] = entry + g.usage = newUsage + size + return nil +} + +// ReleaseTree forgets accounting for files at or below path. Recovery cleanup +// can call this after it successfully removes a stale request tree. +func (g *EphemeralCapacityGuard) ReleaseTree(path string) error { + cleanPath, _, err := g.registeredPath(path) + if err != nil { + return err + } + + g.mu.Lock() + defer g.mu.Unlock() + for g.hasOpenWriterLocked(cleanPath) { + g.changed.Wait() + } + for entryPath := range g.entries { + if ephemeralPathAtOrBelow(entryPath, cleanPath) { + g.releaseLocked(entryPath) + } + } + return nil +} + +// HasActiveReservation reports whether path itself or a descendant has an +// active reservation. Recovery cleanup uses it to avoid active request trees. +func (g *EphemeralCapacityGuard) HasActiveReservation(path string) bool { + cleanPath, _, err := g.registeredPath(path) + if err != nil { + return false + } + + g.mu.Lock() + defer g.mu.Unlock() + for entryPath, entry := range g.entries { + if entry.isActive() && ephemeralPathAtOrBelow(entryPath, cleanPath) { + return true + } + } + return false +} + +func (g *EphemeralCapacityGuard) commitWaiterCount(path string) int { + cleanPath, _, err := g.registeredPath(path) + if err != nil { + return 0 + } + + g.mu.Lock() + defer g.mu.Unlock() + return g.commitWaiters[cleanPath] +} + +// NewWriter returns a writer that admits unknown-length input in bounded +// chunks. Callers must close it before committing or releasing the path. +func (g *EphemeralCapacityGuard) NewWriter(path string, destination io.Writer) (*EphemeralCapacityWriter, error) { + if destination == nil { + return nil, fmt.Errorf("ephemeral writer destination is nil") + } + cleanPath, root, err := g.registeredPath(path) + if err != nil { + return nil, err + } + + g.mu.Lock() + defer g.mu.Unlock() + entry, found := g.entries[cleanPath] + if !found || !entry.isActive() { + if err := g.checkCapacityLocked(root, 0); err != nil { + return nil, err + } + entry.state = ephemeralCapacityActive + entry.reserved = 0 + entry.pending = 0 + entry.inflight = 0 + } + entry.state = ephemeralCapacityWriting + entry.openWriters++ + g.entries[cleanPath] = entry + return &EphemeralCapacityWriter{guard: g, path: cleanPath, destination: destination}, nil +} + +// EphemeralCapacityWriter bounds writes through an EphemeralCapacityGuard. +// Close finalizes its accounting lifecycle without closing the destination. +type EphemeralCapacityWriter struct { + mu sync.Mutex + guard *EphemeralCapacityGuard + path string + destination io.Writer + closed bool +} + +var _ io.WriteCloser = (*EphemeralCapacityWriter)(nil) + +func (w *EphemeralCapacityWriter) Write(payload []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + if w.closed { + return 0, fmt.Errorf("ephemeral capacity writer is closed") + } + + total := 0 + for len(payload) > 0 { + chunkLength := min(len(payload), int(ephemeralCapacityWriteChunk)) + grown, err := w.guard.beginWrite(w.path, int64(chunkLength)) + if err != nil { + return total, err + } + written, writeErr := w.destination.Write(payload[:chunkLength]) + if written < 0 || written > chunkLength { + w.guard.settleWrite(w.path, int64(chunkLength), 0, grown) + return total, fmt.Errorf("ephemeral destination returned invalid write count %d", written) + } + w.guard.settleWrite(w.path, int64(chunkLength), int64(written), grown) + total += written + payload = payload[written:] + if writeErr != nil { + return total, writeErr + } + if written != chunkLength { + return total, io.ErrShortWrite + } + } + return total, nil +} + +func (w *EphemeralCapacityWriter) Close() error { + w.mu.Lock() + defer w.mu.Unlock() + if w.closed { + return nil + } + w.closed = true + w.guard.closeWriter(w.path) + return nil +} + +func (e ephemeralCapacityEntry) isActive() bool { + return e.state == ephemeralCapacityActive || e.state == ephemeralCapacityWriting +} + +func (g *EphemeralCapacityGuard) accountExistingFiles(root string) error { + err := filepath.WalkDir(root, func(path string, dirEntry os.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if dirEntry.Type()&os.ModeSymlink != 0 { + if dirEntry.IsDir() { + return filepath.SkipDir + } + return nil + } + info, err := dirEntry.Info() + if err != nil { + return err + } + if !info.Mode().IsRegular() { + return nil + } + cleanPath := filepath.Clean(path) + if _, found := g.entries[cleanPath]; found { + return nil + } + if info.Size() > math.MaxInt64-g.usage { + return fmt.Errorf("ephemeral usage exceeds supported size") + } + g.entries[cleanPath] = ephemeralCapacityEntry{ + state: ephemeralCapacityExisting, baseline: info.Size(), + } + g.usage += info.Size() + return nil + }) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return fmt.Errorf("accounting ephemeral root %q: %w", root, err) + } + return nil +} + +func (g *EphemeralCapacityGuard) beginWrite(path string, size int64) (int64, error) { + _, root, err := g.registeredPath(path) + if err != nil { + return 0, err + } + + g.mu.Lock() + defer g.mu.Unlock() + entry, found := g.entries[path] + if !found || entry.state != ephemeralCapacityWriting || entry.openWriters == 0 { + return 0, fmt.Errorf("ephemeral path %q has no active writer", path) + } + availableReservation := entry.pending - entry.inflight + grow := max(int64(0), size-availableReservation) + if err := g.checkCapacityLocked(root, grow); err != nil { + return 0, err + } + entry.reserved += grow + entry.pending += grow + entry.inflight += size + g.entries[path] = entry + g.usage += grow + return grow, nil +} + +func (g *EphemeralCapacityGuard) settleWrite(path string, attempted, written, grown int64) { + g.mu.Lock() + defer g.mu.Unlock() + entry, found := g.entries[path] + if !found { + return + } + entry.inflight -= attempted + entry.pending -= written + unusedGrowth := min(grown, attempted-written) + entry.reserved -= unusedGrowth + entry.pending -= unusedGrowth + g.usage -= unusedGrowth + g.entries[path] = entry +} + +func (g *EphemeralCapacityGuard) closeWriter(path string) { + g.mu.Lock() + defer g.mu.Unlock() + entry, found := g.entries[path] + if !found || entry.openWriters == 0 { + return + } + entry.openWriters-- + if entry.openWriters == 0 { + entry.state = ephemeralCapacityActive + } + g.entries[path] = entry + g.changed.Broadcast() +} + +func (g *EphemeralCapacityGuard) checkCapacityLocked(root string, requested int64) error { + available, err := ephemeralAvailableBytes(root) + if err != nil { + return err + } + exceedsLimit := requested > g.byteLimit-g.usage + pending := g.pendingLocked() + freeAfterHeadroom := available - min(available, g.minFreeBytes) + exceedsAvailable := pending > freeAfterHeadroom || requested > freeAfterHeadroom-pending + if exceedsLimit || exceedsAvailable { + return &EphemeralCapacityError{ + RequestedBytes: requested, + UsageBytes: g.usage, + LimitBytes: g.byteLimit, + AvailableBytes: available, + HeadroomBytes: g.minFreeBytes, + } + } + return nil +} + +func (g *EphemeralCapacityGuard) pendingLocked() int64 { + var pending int64 + for _, entry := range g.entries { + if entry.pending > math.MaxInt64-pending { + return math.MaxInt64 + } + pending += entry.pending + } + return pending +} + +func (g *EphemeralCapacityGuard) hasOpenWriterLocked(path string) bool { + for entryPath, entry := range g.entries { + if entry.openWriters > 0 && ephemeralPathAtOrBelow(entryPath, path) { + return true + } + } + return false +} + +func (g *EphemeralCapacityGuard) releaseLocked(path string) { + entry, found := g.entries[path] + if !found { + return + } + g.usage -= entry.baseline + entry.reserved + delete(g.entries, path) +} + +func (g *EphemeralCapacityGuard) registeredPath(path string) (string, string, error) { + cleanPath, err := cleanEphemeralAbsolutePath(path) + if err != nil { + return "", "", err + } + root := "" + for _, candidate := range g.roots { + if ephemeralPathAtOrBelow(cleanPath, candidate) && len(candidate) > len(root) { + root = candidate + } + } + if root == "" { + return "", "", fmt.Errorf("path %q is outside registered ephemeral roots", cleanPath) + } + if err := rejectEphemeralSymlinkComponents(cleanPath); err != nil { + return "", "", fmt.Errorf("validating ephemeral path %q: %w", cleanPath, err) + } + return cleanPath, root, nil +} + +func cleanEphemeralAbsolutePath(path string) (string, error) { + if path == "" { + return "", fmt.Errorf("ephemeral path is empty") + } + absPath, err := filepath.Abs(path) + if err != nil { + return "", fmt.Errorf("resolving ephemeral path %q: %w", path, err) + } + return filepath.Clean(absPath), nil +} + +func rejectEphemeralSymlinkComponents(path string) error { + volume := filepath.VolumeName(path) + remainder := strings.TrimPrefix(path, volume) + current := volume + string(filepath.Separator) + remainder = strings.TrimPrefix(remainder, string(filepath.Separator)) + for _, component := range strings.Split(remainder, string(filepath.Separator)) { + if component == "" { + continue + } + current = filepath.Join(current, component) + info, err := os.Lstat(current) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return fmt.Errorf("checking path component %q: %w", current, err) + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("path component %q is a symlink", current) + } + } + return nil +} + +func ephemeralPathAtOrBelow(path, parent string) bool { + return path == parent || strings.HasPrefix(path, parent+string(filepath.Separator)) +} + +func ephemeralAvailableBytes(path string) (int64, error) { + diskInfo, err := xsysinfo.GetDiskInfo(path) + if err != nil { + return 0, fmt.Errorf("reading ephemeral filesystem availability: %w", err) + } + if diskInfo.Available > math.MaxInt64 { + return math.MaxInt64, nil + } + return int64(diskInfo.Available), nil +} diff --git a/core/services/worker/ephemeral_capacity_test.go b/core/services/worker/ephemeral_capacity_test.go new file mode 100644 index 000000000..41a44db93 --- /dev/null +++ b/core/services/worker/ephemeral_capacity_test.go @@ -0,0 +1,317 @@ +// SPDX-License-Identifier: MIT + +package worker + +import ( + "bytes" + "errors" + "io" + "os" + "path/filepath" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type capacityGatedFileWriter struct { + file *os.File + entered chan struct{} + resume chan struct{} + once sync.Once +} + +func (w *capacityGatedFileWriter) Write(p []byte) (int, error) { + w.once.Do(func() { + close(w.entered) + <-w.resume + }) + return w.file.Write(p) +} + +type capacityShortWriter struct{} + +func (capacityShortWriter) Write(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + return 1, nil +} + +var _ = Describe("EphemeralCapacityGuard", func() { + It("accounts existing regular files without following symlinks", func() { + root := GinkgoT().TempDir() + outside := filepath.Join(GinkgoT().TempDir(), "outside.bin") + Expect(os.WriteFile(filepath.Join(root, "existing.bin"), make([]byte, 6), 0o600)).To(Succeed()) + Expect(os.WriteFile(outside, make([]byte, 100), 0o600)).To(Succeed()) + Expect(os.Symlink(outside, filepath.Join(root, "outside-link"))).To(Succeed()) + + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + + err = guard.Reserve(filepath.Join(root, "next.bin"), 5) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(capacityErr.RequestedBytes).To(Equal(int64(5))) + Expect(capacityErr.UsageBytes).To(Equal(int64(6))) + Expect(capacityErr.LimitBytes).To(Equal(int64(10))) + Expect(capacityErr.AvailableBytes).To(BeNumerically(">", 0)) + Expect(capacityErr.HeadroomBytes).To(BeZero()) + }) + + It("serializes competing reservations", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 1, 0) + Expect(err).NotTo(HaveOccurred()) + + start := make(chan struct{}) + results := make(chan error, 32) + var wait sync.WaitGroup + for i := range 32 { + wait.Add(1) + go func(index int) { + defer wait.Done() + <-start + results <- guard.Reserve(filepath.Join(root, string(rune('a'+index))), 1) + }(i) + } + close(start) + wait.Wait() + close(results) + + succeeded := 0 + for result := range results { + if result == nil { + succeeded++ + continue + } + var capacityErr *EphemeralCapacityError + Expect(errors.As(result, &capacityErr)).To(BeTrue()) + } + Expect(succeeded).To(Equal(1)) + }) + + It("makes only an equal active reservation idempotent", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + path := filepath.Join(root, "nested", "payload.bin") + + Expect(guard.Reserve(path, 4)).To(Succeed()) + Expect(guard.Reserve(filepath.Join(root, "nested", ".", "payload.bin"), 4)).To(Succeed()) + err = guard.Reserve(path, 5) + var conflictErr *EphemeralReservationConflictError + Expect(errors.As(err, &conflictErr)).To(BeTrue()) + Expect(conflictErr.ActiveBytes).To(Equal(int64(4))) + Expect(conflictErr.RequestedBytes).To(Equal(int64(5))) + Expect(guard.Reserve(filepath.Join(root, "other.bin"), 6)).To(Succeed()) + Expect(guard.Release(filepath.Join(root, "nested", ".", "payload.bin"))).To(Succeed()) + Expect(guard.Release(path)).To(Succeed()) + Expect(guard.Reserve(filepath.Join(root, "replacement.bin"), 4)).To(Succeed()) + }) + + It("retains committed bytes when the same path starts another reservation", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + path := filepath.Join(root, "payload.bin") + + Expect(guard.Reserve(path, 4)).To(Succeed()) + Expect(os.WriteFile(path, make([]byte, 4), 0o600)).To(Succeed()) + Expect(guard.Commit(path)).To(Succeed()) + Expect(guard.HasActiveReservation(path)).To(BeFalse()) + Expect(guard.Reserve(path, 6)).To(Succeed()) + Expect(guard.HasActiveReservation(path)).To(BeTrue()) + + err = guard.Reserve(filepath.Join(root, "overflow.bin"), 1) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(capacityErr.UsageBytes).To(Equal(int64(10))) + }) + + It("retains startup-accounted bytes when the path is reserved", func() { + root := GinkgoT().TempDir() + path := filepath.Join(root, "payload.bin") + Expect(os.WriteFile(path, make([]byte, 4), 0o600)).To(Succeed()) + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + + Expect(guard.Reserve(path, 6)).To(Succeed()) + err = guard.Reserve(filepath.Join(root, "overflow.bin"), 1) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(capacityErr.UsageBytes).To(Equal(int64(10))) + }) + + It("commits the regular file's actual size", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + path := filepath.Join(root, "payload.bin") + + Expect(guard.Reserve(path, 10)).To(Succeed()) + Expect(os.WriteFile(path, []byte("four"), 0o600)).To(Succeed()) + Expect(guard.Commit(path)).To(Succeed()) + Expect(guard.Commit(filepath.Clean(path))).To(Succeed()) + Expect(guard.Reserve(filepath.Join(root, "six.bin"), 6)).To(Succeed()) + }) + + It("preserves configured filesystem headroom", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 1<<30, 1<<62) + Expect(err).NotTo(HaveOccurred()) + + err = guard.Reserve(filepath.Join(root, "payload.bin"), 1) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(capacityErr.RequestedBytes).To(Equal(int64(1))) + Expect(capacityErr.AvailableBytes).To(BeNumerically(">", 0)) + Expect(capacityErr.HeadroomBytes).To(Equal(int64(1 << 62))) + }) + + It("reserves bounded chunks before forwarding unknown-length input", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, ephemeralCapacityWriteChunk+1, 0) + Expect(err).NotTo(HaveOccurred()) + path := filepath.Join(root, "payload.bin") + file, err := os.Create(path) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(file.Close) + writer, err := guard.NewWriter(path, file) + Expect(err).NotTo(HaveOccurred()) + + n, err := writer.Write(make([]byte, ephemeralCapacityWriteChunk+2)) + Expect(n).To(Equal(int(ephemeralCapacityWriteChunk))) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(writer.Close()).To(Succeed()) + info, err := file.Stat() + Expect(err).NotTo(HaveOccurred()) + Expect(info.Size()).To(Equal(ephemeralCapacityWriteChunk)) + }) + + It("waits for an open bounded writer before committing", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + path := filepath.Join(root, "payload.bin") + file, err := os.Create(path) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(file.Close) + gated := &capacityGatedFileWriter{ + file: file, entered: make(chan struct{}), resume: make(chan struct{}), + } + DeferCleanup(func() { + select { + case <-gated.resume: + default: + close(gated.resume) + } + }) + writer, err := guard.NewWriter(path, gated) + Expect(err).NotTo(HaveOccurred()) + + writeDone := make(chan error, 1) + go func() { + _, writeErr := writer.Write([]byte("1234567")) + writeDone <- writeErr + }() + Eventually(gated.entered).Should(BeClosed()) + commitDone := make(chan error, 1) + go func() { commitDone <- guard.Commit(path) }() + Eventually(func() int { return guard.commitWaiterCount(path) }).Should(Equal(1)) + Expect(commitDone).NotTo(Receive()) + + close(gated.resume) + Expect(<-writeDone).To(Succeed()) + Expect(commitDone).NotTo(Receive()) + Expect(writer.Close()).To(Succeed()) + Eventually(commitDone).Should(Receive(Succeed())) + + err = guard.Reserve(filepath.Join(root, "other.bin"), 4) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(capacityErr.UsageBytes).To(Equal(int64(7))) + }) + + It("does not share pending capacity between concurrent writers", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + path := filepath.Join(root, "payload.bin") + file, err := os.Create(path) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(file.Close) + gated := &capacityGatedFileWriter{ + file: file, entered: make(chan struct{}), resume: make(chan struct{}), + } + first, err := guard.NewWriter(path, gated) + Expect(err).NotTo(HaveOccurred()) + var secondDestination bytes.Buffer + second, err := guard.NewWriter(path, &secondDestination) + Expect(err).NotTo(HaveOccurred()) + + firstDone := make(chan error, 1) + go func() { + _, writeErr := first.Write([]byte("1234567")) + firstDone <- writeErr + }() + Eventually(gated.entered).Should(BeClosed()) + + n, err := second.Write([]byte("7654321")) + Expect(n).To(BeZero()) + var capacityErr *EphemeralCapacityError + Expect(errors.As(err, &capacityErr)).To(BeTrue()) + Expect(secondDestination.Len()).To(BeZero()) + + close(gated.resume) + Expect(<-firstDone).To(Succeed()) + Expect(first.Close()).To(Succeed()) + Expect(second.Close()).To(Succeed()) + }) + + It("rolls back bytes the destination writer does not accept", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 5, 0) + Expect(err).NotTo(HaveOccurred()) + writer, err := guard.NewWriter(filepath.Join(root, "payload.bin"), capacityShortWriter{}) + Expect(err).NotTo(HaveOccurred()) + + n, err := writer.Write([]byte("123")) + Expect(n).To(Equal(1)) + Expect(err).To(MatchError(io.ErrShortWrite)) + Expect(writer.Close()).To(Succeed()) + Expect(guard.Reserve(filepath.Join(root, "other.bin"), 4)).To(Succeed()) + }) + + It("rejects paths outside roots and through symlinks", func() { + root := GinkgoT().TempDir() + outside := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 100, 0) + Expect(err).NotTo(HaveOccurred()) + + Expect(guard.Reserve(filepath.Join(outside, "payload.bin"), 1)).To( + MatchError(ContainSubstring("outside registered ephemeral roots")), + ) + Expect(os.Symlink(outside, filepath.Join(root, "escape"))).To(Succeed()) + Expect(guard.Reserve(filepath.Join(root, "escape", "payload.bin"), 1)).To( + MatchError(ContainSubstring("symlink")), + ) + }) + + It("supports recovery tree accounting without dropping active reservations", func() { + root := GinkgoT().TempDir() + guard, err := NewEphemeralCapacityGuard([]string{root}, 10, 0) + Expect(err).NotTo(HaveOccurred()) + active := filepath.Join(root, "active", "payload.bin") + stale := filepath.Join(root, "stale", "payload.bin") + + Expect(guard.Reserve(active, 4)).To(Succeed()) + Expect(guard.Account(stale, 3)).To(Succeed()) + Expect(guard.HasActiveReservation(root)).To(BeTrue()) + Expect(guard.ReleaseTree(filepath.Join(root, "stale"))).To(Succeed()) + Expect(guard.Reserve(filepath.Join(root, "replacement.bin"), 6)).To(Succeed()) + Expect(guard.ReleaseTree(root)).To(Succeed()) + Expect(guard.HasActiveReservation(root)).To(BeFalse()) + }) +})