diff --git a/core/services/nodes/file_stager_release_test.go b/core/services/nodes/file_stager_release_test.go index bafe6c8f6..4796ff5b9 100644 --- a/core/services/nodes/file_stager_release_test.go +++ b/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)) + }) }) diff --git a/core/services/nodes/file_stager_s3.go b/core/services/nodes/file_stager_s3.go index e63e0a629..707d0cc87 100644 --- a/core/services/nodes/file_stager_s3.go +++ b/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 != "" { diff --git a/core/services/nodes/file_transfer_server.go b/core/services/nodes/file_transfer_server.go index 42fcc3a02..9e26c9b17 100644 --- a/core/services/nodes/file_transfer_server.go +++ b/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)) diff --git a/core/services/nodes/file_transfer_server_test.go b/core/services/nodes/file_transfer_server_test.go index 03912a965..164019b7a 100644 --- a/core/services/nodes/file_transfer_server_test.go +++ b/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)