fix(distributed): close staging accounting gaps

Keep unknown-length reservations charged until bytes reach disk and bound NATS release waits by the lifecycle cleanup deadline.

Assisted-by: Codex:gpt-6

core/services/nodes/file_stager_release_test.go

@@ -24,6 +24,7 @@ type releaseTestMessaging struct {
 	payload       []byte
 	onRequest     func()
 	requestCalled bool
+	timeout       time.Duration
 }
 
 func (m *releaseTestMessaging) Publish(string, any) error { return nil }
@@ -39,10 +40,11 @@ func (m *releaseTestMessaging) QueueSubscribeReply(string, string, func([]byte,
 func (m *releaseTestMessaging) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
 	return releaseTestSubscription{}, nil
 }
-func (m *releaseTestMessaging) Request(subject string, data []byte, _ time.Duration) ([]byte, error) {
+func (m *releaseTestMessaging) Request(subject string, data []byte, timeout time.Duration) ([]byte, error) {
 	m.subject = subject
 	m.payload = append([]byte(nil), data...)
 	m.requestCalled = true
+	m.timeout = timeout
 	if m.onRequest != nil {
 		m.onRequest()
 	}
@@ -191,4 +193,33 @@ var _ = Describe("File stager exact-key release", func() {
 		Expect(err).NotTo(HaveOccurred())
 		Expect(exists).To(BeFalse())
 	})
+
+	It("does not send a release request after cleanup is canceled", func() {
+		store, err := storage.NewFilesystemStore(GinkgoT().TempDir())
+		Expect(err).NotTo(HaveOccurred())
+		fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
+		Expect(err).NotTo(HaveOccurred())
+		client := &releaseTestMessaging{}
+		stager := NewS3NATSFileStager(fm, client)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		Expect(stager.ReleaseRemote(ctx, "node.one", "ephemeral/request-id/audio/input.wav")).To(MatchError(context.Canceled))
+		Expect(client.requestCalled).To(BeFalse())
+	})
+
+	It("bounds the NATS release wait by the remaining cleanup deadline", func() {
+		store, err := storage.NewFilesystemStore(GinkgoT().TempDir())
+		Expect(err).NotTo(HaveOccurred())
+		fm, err := storage.NewFileManager(store, GinkgoT().TempDir())
+		Expect(err).NotTo(HaveOccurred())
+		client := &releaseTestMessaging{}
+		stager := NewS3NATSFileStager(fm, client)
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		Expect(stager.ReleaseRemote(ctx, "node.one", "ephemeral/request-id/audio/input.wav")).To(Succeed())
+		Expect(client.timeout).To(BeNumerically(">", time.Second))
+		Expect(client.timeout).To(BeNumerically("<=", 2*time.Second))
+	})
 })

core/services/nodes/file_stager_s3.go

@@ -196,13 +196,30 @@ func (s *S3NATSFileStager) ReleaseRemote(ctx context.Context, nodeID, key string
 	if err := validateEphemeralReleaseKey(key); err != nil {
 		return err
 	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	timeout := 30 * time.Second
+	if deadline, ok := ctx.Deadline(); ok {
+		remaining := time.Until(deadline)
+		if remaining <= 0 {
+			return context.DeadlineExceeded
+		}
+		timeout = min(timeout, remaining)
+	}
 	reply, err := messaging.RequestJSON[fileReleaseRequest, fileReleaseReply](
 		s.nats,
 		messaging.SubjectNodeFilesRelease(nodeID),
 		fileReleaseRequest{Key: key},
-		30*time.Second,
+		timeout,
 	)
 	if err != nil {
+		if contextErr := ctx.Err(); contextErr != nil {
+			return contextErr
+		}
+		return err
+	}
+	if err := ctx.Err(); err != nil {
 		return err
 	}
 	if reply.Error != "" {

core/services/nodes/file_transfer_server.go

@@ -358,29 +358,26 @@ func (w *uploadStatusWriter) Write(payload []byte) (int, error) {
 	return w.ResponseWriter.Write(payload)
 }
 
-type capacityRequestBody struct {
-	io.ReadCloser
-	writer io.WriteCloser
-}
-
 type ephemeralCapacityWriteError struct{ err error }
 
 func (e *ephemeralCapacityWriteError) Error() string { return e.err.Error() }
 func (e *ephemeralCapacityWriteError) Unwrap() error { return e.err }
 
-func (r *capacityRequestBody) Read(payload []byte) (int, error) {
-	n, readErr := r.ReadCloser.Read(payload)
-	if n == 0 {
-		return n, readErr
+type ephemeralCapacityWriteCloser struct{ io.WriteCloser }
+
+func (w ephemeralCapacityWriteCloser) Write(payload []byte) (int, error) {
+	written, err := w.WriteCloser.Write(payload)
+	if err != nil {
+		return written, &ephemeralCapacityWriteError{err: err}
 	}
-	written, writeErr := r.writer.Write(payload[:n])
-	if writeErr != nil {
-		return written, &ephemeralCapacityWriteError{err: writeErr}
+	return written, nil
+}
+
+func (w ephemeralCapacityWriteCloser) Close() error {
+	if err := w.WriteCloser.Close(); err != nil {
+		return &ephemeralCapacityWriteError{err: err}
 	}
-	if written != n {
-		return written, io.ErrShortWrite
-	}
-	return n, readErr
+	return nil
 }
 
 func uploadWriteStatus(err error) int {
@@ -474,8 +471,8 @@ func handleUploadWithCapacity(w http.ResponseWriter, r *http.Request, stagingDir
 		return
 	}
 
-	var capacityBody *capacityRequestBody
 	capacityEnabled := capacity != nil && targetDir == stagingDir && strings.HasPrefix(key, "ephemeral/")
+	unknownLengthCapacity := capacityEnabled && r.ContentLength < 0
 	capacityPaths := []string{dstPath, dstPath + hashSidecarSuffix, dstPath + targetSidecarSuffix}
 	if capacityEnabled {
 		if r.ContentLength >= 0 {
@@ -483,20 +480,9 @@ func handleUploadWithCapacity(w http.ResponseWriter, r *http.Request, stagingDir
 				http.Error(w, err.Error(), http.StatusInsufficientStorage)
 				return
 			}
-		} else {
-			writer, err := capacity.CapacityWriter(dstPath, io.Discard)
-			if err != nil {
-				http.Error(w, err.Error(), http.StatusInsufficientStorage)
-				return
-			}
-			capacityBody = &capacityRequestBody{ReadCloser: r.Body, writer: writer}
-			r.Body = capacityBody
 		}
 		for _, sidecarPath := range capacityPaths[1:] {
 			if err := capacity.Reserve(sidecarPath, sha256.Size*2); err != nil {
-				if capacityBody != nil {
-					_ = capacityBody.writer.Close()
-				}
 				for _, reservedPath := range capacityPaths {
 					reconcileEphemeralCapacity(capacity, reservedPath, 0)
 				}
@@ -507,22 +493,21 @@ func handleUploadWithCapacity(w http.ResponseWriter, r *http.Request, stagingDir
 	}
 
 	statusWriter := &uploadStatusWriter{ResponseWriter: w}
+	var uploadCapacity EphemeralCapacity
+	if unknownLengthCapacity {
+		uploadCapacity = capacity
+	}
 
 	if cr == nil {
 		// Non-resumable (legacy) path: truncate-create, single fire-and-forget.
-		handleFullUpload(statusWriter, r, dstPath, key, expectedFinalHash)
+		handleFullUpload(statusWriter, r, dstPath, key, expectedFinalHash, uploadCapacity)
 	} else {
-		handleRangeUpload(statusWriter, r, dstPath, key, cr, expectedFinalHash)
+		handleRangeUpload(statusWriter, r, dstPath, key, cr, expectedFinalHash, uploadCapacity)
 	}
 
 	if !capacityEnabled {
 		return
 	}
-	if capacityBody != nil {
-		if err := capacityBody.writer.Close(); err != nil {
-			xlog.Warn("Closing ephemeral capacity writer failed", "path", dstPath, "error", err)
-		}
-	}
 	for _, capacityPath := range capacityPaths {
 		reconcileEphemeralCapacity(capacity, capacityPath, statusWriter.status)
 	}
@@ -540,7 +525,7 @@ func reconcileEphemeralCapacity(capacity EphemeralCapacity, path string, status
 
 // handleFullUpload writes the entire request body to dstPath, replacing any
 // existing content. This is the legacy happy-path with no Range header.
-func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expectedFinalHash string) {
+func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expectedFinalHash string, capacity EphemeralCapacity) {
 	// Reset any in-progress resumable state.
 	_ = os.Remove(dstPath + targetSidecarSuffix)
 
@@ -551,8 +536,27 @@ func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expe
 	}
 	defer f.Close()
 
+	var destination io.Writer = f
+	var capacityWriter io.WriteCloser
+	if capacity != nil {
+		var writer io.WriteCloser
+		writer, err = capacity.CapacityWriter(dstPath, f)
+		if err != nil {
+			_ = os.Remove(dstPath)
+			http.Error(w, err.Error(), http.StatusInsufficientStorage)
+			return
+		}
+		capacityWriter = ephemeralCapacityWriteCloser{WriteCloser: writer}
+		destination = capacityWriter
+	}
+
 	hasher := sha256.New()
-	n, err := io.Copy(f, io.TeeReader(r.Body, hasher))
+	n, err := io.Copy(destination, io.TeeReader(r.Body, hasher))
+	if capacityWriter != nil {
+		if closeErr := capacityWriter.Close(); err == nil {
+			err = closeErr
+		}
+	}
 	if err != nil {
 		os.Remove(dstPath)
 		os.Remove(dstPath + hashSidecarSuffix)
@@ -586,7 +590,7 @@ func handleFullUpload(w http.ResponseWriter, r *http.Request, dstPath, key, expe
 // the request starts at the current file size. When the slice completes the
 // transfer (end+1 == total), it validates the optional expected final hash and
 // writes the sidecar.
-func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key string, cr *contentRange, expectedFinalHash string) {
+func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key string, cr *contentRange, expectedFinalHash string, capacity EphemeralCapacity) {
 	// Determine the current on-disk size (0 if missing).
 	var currentSize int64
 	if info, err := os.Stat(dstPath); err == nil {
@@ -670,6 +674,18 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
 		return
 	}
 	defer func() { _ = f.Close() }()
+	var destination io.Writer = f
+	var capacityWriter io.WriteCloser
+	if capacity != nil {
+		var writer io.WriteCloser
+		writer, err = capacity.CapacityWriter(dstPath, f)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInsufficientStorage)
+			return
+		}
+		capacityWriter = ephemeralCapacityWriteCloser{WriteCloser: writer}
+		destination = capacityWriter
+	}
 
 	// Persist the declared expected hash so subsequent chunks can be
 	// cross-checked.
@@ -681,7 +697,12 @@ func handleRangeUpload(w http.ResponseWriter, r *http.Request, dstPath, key stri
 
 	expectedChunkLen := cr.end - cr.start + 1
 	limited := io.LimitReader(r.Body, expectedChunkLen)
-	n, err := io.Copy(f, limited)
+	n, err := io.Copy(destination, limited)
+	if capacityWriter != nil {
+		if closeErr := capacityWriter.Close(); err == nil {
+			err = closeErr
+		}
+	}
 	if err != nil {
 		xlog.Error("Range upload chunk failed", "key", key, "bytesReceived", n, "expected", expectedChunkLen, "remote", r.RemoteAddr, "error", err)
 		http.Error(w, fmt.Sprintf("writing file: %v", err), uploadWriteStatus(err))

core/services/nodes/file_transfer_server_test.go

@@ -39,6 +39,38 @@ type failingWriteCloser struct{ err error }
 func (w failingWriteCloser) Write([]byte) (int, error) { return 0, w.err }
 func (failingWriteCloser) Close() error                { return nil }
 
+type blockingDestinationWriteCloser struct {
+	destination io.Writer
+	written     chan struct{}
+	release     chan struct{}
+}
+
+func (w *blockingDestinationWriteCloser) Write(payload []byte) (int, error) {
+	n, err := w.destination.Write(payload)
+	close(w.written)
+	<-w.release
+	return n, err
+}
+
+func (*blockingDestinationWriteCloser) Close() error { return nil }
+
+type blockingDestinationCapacity struct {
+	written chan struct{}
+	release chan struct{}
+}
+
+func (*blockingDestinationCapacity) Reserve(string, int64) error { return nil }
+func (*blockingDestinationCapacity) Commit(string) error         { return nil }
+func (*blockingDestinationCapacity) Claim(string) error          { return nil }
+func (*blockingDestinationCapacity) Release(string) error        { return nil }
+func (g *blockingDestinationCapacity) CapacityWriter(_ string, destination io.Writer) (io.WriteCloser, error) {
+	return &blockingDestinationWriteCloser{
+		destination: destination,
+		written:     g.written,
+		release:     g.release,
+	}, nil
+}
+
 func (g *recordingEphemeralCapacity) Reserve(_ string, size int64) error {
 	g.reserved = size
 	return g.reserveErr
@@ -117,6 +149,40 @@ var _ = Describe("FileTransferServer", func() {
 			Expect(recorder.Code).To(Equal(http.StatusInsufficientStorage))
 		})
 
+		It("keeps unknown-length bytes guarded until they reach the staged file", func() {
+			stagingDir := GinkgoT().TempDir()
+			modelsDir := GinkgoT().TempDir()
+			dataDir := GinkgoT().TempDir()
+			guard := &blockingDestinationCapacity{
+				written: make(chan struct{}),
+				release: make(chan struct{}),
+			}
+			released := false
+			defer func() {
+				if !released {
+					close(guard.release)
+				}
+			}()
+			recorder := httptest.NewRecorder()
+			request := httptest.NewRequest(http.MethodPut, "/v1/files/ephemeral/audio/request/input.wav", strings.NewReader("payload"))
+			request.ContentLength = -1
+			done := make(chan struct{})
+
+			go func() {
+				defer GinkgoRecover()
+				handleUploadWithCapacity(recorder, request, stagingDir, modelsDir, dataDir, "ephemeral/audio/request/input.wav", 0, guard)
+				close(done)
+			}()
+
+			Eventually(guard.written).Should(BeClosed())
+			path := filepath.Join(stagingDir, "ephemeral", "audio", "request", "input.wav")
+			Expect(os.ReadFile(path)).To(Equal([]byte("payload")))
+			close(guard.release)
+			released = true
+			Eventually(done).Should(BeClosed())
+			Expect(recorder.Code).To(Equal(http.StatusOK))
+		})
+
 		It("round-trips file content correctly", func() {
 			ts, _, _, _ := setupTestServer("secret-token", 0)