LocalAI/pkg/downloader/stall_test.go

package downloader_test

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"time"

	. "github.com/mudler/LocalAI/pkg/downloader"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("Download stall timeout", func() {
	var filePath string
	var savedTimeout time.Duration

	BeforeEach(func() {
		dir, err := os.Getwd()
		Expect(err).ToNot(HaveOccurred())
		filePath = dir + "/stall_model"
		savedTimeout = DownloadStallTimeout
	})

	AfterEach(func() {
		DownloadStallTimeout = savedTimeout
		_ = os.Remove(filePath)
		_ = os.Remove(filePath + ".partial")
	})

	It("aborts a download that stalls mid-stream instead of hanging forever", func() {
		// Server sends a chunk, flushes, then blocks forever without closing
		// the connection — a silently-dropped TCP stream. Without a stall
		// guard the body Read blocks indefinitely and DownloadFile never
		// returns.
		release := make(chan struct{})
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method == "HEAD" {
				w.Header().Set("Accept-Ranges", "bytes")
				w.WriteHeader(http.StatusOK)
				return
			}
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write(make([]byte, 4096))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			<-release // hang: no more data, never close
		}))
		defer server.Close()
		defer close(release)

		DownloadStallTimeout = 300 * time.Millisecond

		done := make(chan error, 1)
		go func() {
			done <- URI(server.URL).DownloadFileWithContext(
				context.Background(), filePath, "", 1, 1,
				func(s1, s2, s3 string, f float64) {})
		}()

		var err error
		Eventually(done, "5s").Should(Receive(&err))
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring("stall"))
	})

	It("preserves the .partial file when a download stalls so it can resume", func() {
		release := make(chan struct{})
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method == "HEAD" {
				w.Header().Set("Accept-Ranges", "bytes")
				w.WriteHeader(http.StatusOK)
				return
			}
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write(make([]byte, 4096))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			<-release
		}))
		defer server.Close()
		defer close(release)

		DownloadStallTimeout = 300 * time.Millisecond

		done := make(chan error, 1)
		go func() {
			done <- URI(server.URL).DownloadFileWithContext(
				context.Background(), filePath, "", 1, 1,
				func(s1, s2, s3 string, f float64) {})
		}()
		Eventually(done, "5s").Should(Receive(HaveOccurred()))

		info, statErr := os.Stat(filePath + ".partial")
		Expect(statErr).ToNot(HaveOccurred(), "the .partial must survive a stall so the next attempt can resume")
		Expect(info.Size()).To(BeNumerically(">", 0))
	})

	It("does not count partial-file hashing time against the stall window when resuming", func() {
		// A resumed download re-hashes the existing .partial before the copy
		// loop starts. On slow storage a multi-GB partial takes longer to hash
		// than the stall window, and since no network read happens while
		// hashing, a watchdog armed before the hash aborts a healthy resume.
		// Every retry then re-pays the same hash and fails identically, so the
		// install wedges permanently (observed with a 7.9GB partial on a CIFS
		// models share). The sparse file keeps disk usage at zero while still
		// forcing the hash to churn through every byte.
		const partialSize = int64(2) << 30
		tail := []byte("remaining-bytes")
		total := partialSize + int64(len(tail))

		partial, err := os.OpenFile(filePath+".partial", os.O_CREATE|os.O_WRONLY, 0600)
		Expect(err).ToNot(HaveOccurred())
		Expect(partial.Truncate(partialSize)).To(Succeed())
		Expect(partial.Close()).To(Succeed())

		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method == "HEAD" {
				w.Header().Set("Accept-Ranges", "bytes")
				w.WriteHeader(http.StatusOK)
				return
			}
			if r.Header.Get("Range") == "" {
				// A resume must ask for a range; anything else would re-download
				// (and here, concatenate) the whole body.
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", partialSize, total-1, total))
			w.WriteHeader(http.StatusPartialContent)
			_, _ = w.Write(tail)
		}))
		defer server.Close()

		DownloadStallTimeout = 150 * time.Millisecond

		err = URI(server.URL).DownloadFileWithContext(
			context.Background(), filePath, "", 1, 1,
			func(s1, s2, s3 string, f float64) {})
		Expect(err).ToNot(HaveOccurred(), "a promptly-answered resume must not be aborted by hash time")

		info, statErr := os.Stat(filePath)
		Expect(statErr).ToNot(HaveOccurred())
		Expect(info.Size()).To(Equal(total))
	})

	It("does not abort a slow-but-steady download", func() {
		// One byte every 100ms keeps the idle clock from ever expiring even
		// though the total transfer outlasts the stall timeout.
		payload := make([]byte, 12)
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method == "HEAD" {
				w.Header().Set("Accept-Ranges", "bytes")
				w.WriteHeader(http.StatusOK)
				return
			}
			w.WriteHeader(http.StatusOK)
			f, _ := w.(http.Flusher)
			for i := range payload {
				_, _ = w.Write(payload[i : i+1])
				if f != nil {
					f.Flush()
				}
				time.Sleep(100 * time.Millisecond)
			}
		}))
		defer server.Close()

		DownloadStallTimeout = 300 * time.Millisecond

		err := URI(server.URL).DownloadFileWithContext(
			context.Background(), filePath, "", 1, 1,
			func(s1, s2, s3 string, f float64) {})
		Expect(err).ToNot(HaveOccurred())
	})
})