Bump ioswmr and fix tee for hf

Author: wzshiming

cache.go

@@ -234,7 +234,6 @@ func (m *MirrorHandler) cacheResponse(w http.ResponseWriter, r *http.Request) {
 			}
 		}
 
-		defer tee.Close()
 		tee.ServeHTTP(w, r)
 		return
 	}
@@ -278,7 +277,7 @@ func (m *MirrorHandler) cacheFile(ctx context.Context, sourceFile, cacheFile str
 }
 
 func (m *MirrorHandler) cacheFileDirect(ctx context.Context, sourceFile, cacheFile string) error {
-	resp, info, err := httpGet(ctx, m.client(), sourceFile)
+	resp, info, err := httpGet(ctx, m.client(), sourceFile, false)
 	if err != nil {
 		return err
 	}

fetch.go

@@ -40,7 +40,7 @@ func httpHead(ctx context.Context, client *http.Client, p string) (fs.FileInfo,
 //
 // The caller is responsible for closing the returned io.ReadCloser.
 // Returns ErrNotOK if the response status is not 200 OK.
-func httpGet(ctx context.Context, client *http.Client, p string) (io.ReadCloser, *fileInfo, error) {
+func httpGet(ctx context.Context, client *http.Client, p string, teeHf bool) (io.ReadCloser, *fileInfo, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, p, nil)
 	if err != nil {
 		return nil, nil, err
@@ -56,7 +56,22 @@ func httpGet(ctx context.Context, client *http.Client, p string) (io.ReadCloser,
 		return nil, nil, fmt.Errorf("%w: http status %d", ErrNotOK, resp.StatusCode)
 	}
 
-	return resp.Body, &fileInfo{
+	body := resp.Body
+	// huggingface HEAD and GET headers may differ, so we perform a HEAD request to get accurate metadata for caching decisions.
+	if teeHf && resp.ContentLength <= 0 {
+		if _, ok := hfHosts[req.URL.Host]; ok {
+			reqHead, err := http.NewRequestWithContext(ctx, http.MethodHead, p, nil)
+			if err != nil {
+				return nil, nil, err
+			}
+			resp, err = client.Do(reqHead)
+			if err != nil {
+				return nil, nil, err
+			}
+		}
+	}
+
+	return body, &fileInfo{
 		name: p,
 		resp: resp,
 	}, nil

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/OpenCIDN/cidn v0.0.108
 	github.com/spf13/pflag v1.0.10
 	github.com/wzshiming/httpseek v0.5.0
-	github.com/wzshiming/ioswmr v0.0.0-20260228060639-c732d91a1feb
+	github.com/wzshiming/ioswmr v0.0.0-20260302055634-59c8070e7d03
 	github.com/wzshiming/sss v0.7.0
 	golang.org/x/sync v0.19.0
 	k8s.io/apimachinery v0.34.1

go.sum

@@ -89,6 +89,12 @@ github.com/wzshiming/httpseek v0.5.0 h1:9pFLlTcebylAtLA4ni0Qo5xwXTgeQOo0iVQA3G1Q
 github.com/wzshiming/httpseek v0.5.0/go.mod h1:YoZhlLIwNjTBDXIT8NpK5zRjOgZouRXPaBfjVXdqMMs=
 github.com/wzshiming/ioswmr v0.0.0-20260228060639-c732d91a1feb h1:XvBxo2V5vkklUAdpC9IhTck8N8EAgrRfRhOwDsXaLcM=
 github.com/wzshiming/ioswmr v0.0.0-20260228060639-c732d91a1feb/go.mod h1:TwwDyS1wnJG3AvKliA+PPB0kliN3yEjsabtH4o7xySQ=
+github.com/wzshiming/ioswmr v0.0.0-20260302052319-c9316c655ea7 h1:9TstNh3ONYZgvJofkcb4IW5h47k1QqKQ2jZisw9S2n0=
+github.com/wzshiming/ioswmr v0.0.0-20260302052319-c9316c655ea7/go.mod h1:TwwDyS1wnJG3AvKliA+PPB0kliN3yEjsabtH4o7xySQ=
+github.com/wzshiming/ioswmr v0.0.0-20260302053347-fc8ba863812e h1:LvTADj966y6pJ8A42Yxxog8BtoMZlxgCdtK2rGHGWBI=
+github.com/wzshiming/ioswmr v0.0.0-20260302053347-fc8ba863812e/go.mod h1:TwwDyS1wnJG3AvKliA+PPB0kliN3yEjsabtH4o7xySQ=
+github.com/wzshiming/ioswmr v0.0.0-20260302055634-59c8070e7d03 h1:rOFrtfMWksCMFwE8IhbCAxkPeBGLy4mOtxrBlhrNhhI=
+github.com/wzshiming/ioswmr v0.0.0-20260302055634-59c8070e7d03/go.mod h1:TwwDyS1wnJG3AvKliA+PPB0kliN3yEjsabtH4o7xySQ=
 github.com/wzshiming/sss v0.7.0 h1:YNGJMJ+LBv9dhEUoWTRRl35tru8V2FoW816AR8fZtLg=
 github.com/wzshiming/sss v0.7.0/go.mod h1:4FNGdjjLOzcmZtcCOW+cpL4dmJvg4pMeek0UQy5/FCg=
 github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=

tee_response.go

@@ -7,20 +7,15 @@ import (
 	"io"
 	"io/fs"
 	"net/http"
-	"os"
 	"path"
-	"sync"
 
 	"github.com/wzshiming/ioswmr"
 )
 
 type teeResponse struct {
-	fileInfo  fs.FileInfo
-	swmr      ioswmr.SWMR
-	tmp       *os.File
-	teeCache  *sync.Map
-	etag      string
-	cacheFile string
+	fileInfo fs.FileInfo
+	swmr     ioswmr.SWMR
+	etag     string
 }
 
 func (t *teeResponse) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -31,91 +26,79 @@ func (t *teeResponse) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		rs := t.swmr.NewReadSeeker(0, int(size))
 		defer rs.Close()
 		name := path.Base(r.URL.Path)
-		w.Header().Set("ETag", t.etag)
+		if t.etag != "" {
+			w.Header().Set("ETag", t.etag)
+		}
 		w.Header().Set("Content-Type", "application/octet-stream")
-		w.Header().Set("Content-Length", fmt.Sprint(size))
 		http.ServeContent(w, r, name, t.fileInfo.ModTime(), rs)
 	} else {
 		rs := t.swmr.NewReader(0)
 		defer rs.Close()
-		w.Header().Set("ETag", t.etag)
-		w.Header().Set("Content-Type", "application/octet-stream")
+		if t.etag != "" {
+			w.Header().Set("ETag", t.etag)
+		}
 		w.WriteHeader(http.StatusOK)
 		if r.Method == http.MethodGet {
 			_, _ = io.Copy(w, rs)
 		}
 	}
 }
 
-func (t *teeResponse) Close() error {
-	if !t.swmr.IsClosed() {
-		return nil
-	}
-	if t.swmr.Using() != 0 {
-		return nil
-	}
-	t.teeCache.Delete(t.cacheFile)
-	err := t.tmp.Close()
-	if err != nil {
-		return err
-	}
-	_ = os.Remove(t.tmp.Name())
-	return nil
-}
-
 func (m *MirrorHandler) cacheFileTee(ctx context.Context, sourceFile, cacheFile string) (*teeResponse, error) {
-	resp, info, err := httpGet(ctx, m.client(), sourceFile)
+	body, info, err := httpGet(ctx, m.client(), sourceFile, true)
 	if err != nil {
 		return nil, err
 	}
 
-	var body io.Reader = resp
-
 	contentLength := info.Size()
 	if contentLength == 0 {
-		_ = resp.Close()
+		_ = body.Close()
 		return nil, ErrNotOK
 	}
 
 	if m.Logger != nil {
 		m.Logger.Println("Tee Cache", cacheFile, contentLength)
 	}
 
-	tmp, err := os.CreateTemp("", "mirror-tee-*")
-	if err != nil {
-		_ = resp.Close()
-		return nil, err
-	}
 	fw, err := m.RemoteCache.Writer(ctx, cacheFile)
 	if err != nil {
 		if m.Logger != nil {
 			m.Logger.Println("Cache writer error", cacheFile, contentLength, err)
 		}
-		_ = resp.Close()
-		_ = tmp.Close()
-		_ = os.Remove(tmp.Name())
+		_ = body.Close()
 		return nil, err
 	}
 
-	swmr := ioswmr.NewSWMR(tmp)
+	swmr := ioswmr.NewSWMR(
+		ioswmr.NewMemoryOrTemporaryFileBuffer(nil, nil),
+		ioswmr.WithAutoClose(),
+		ioswmr.WithBeforeCloseFunc(func() {
+			m.teeCache.Delete(cacheFile)
+			if m.Logger != nil {
+				m.Logger.Println("Tee Cache closed", cacheFile, err)
+			}
+		}),
+	)
 
 	tee := &teeResponse{
-		fileInfo:  info,
-		swmr:      swmr,
-		tmp:       tmp,
-		etag:      info.ETag(),
-		teeCache:  &m.teeCache,
-		cacheFile: cacheFile,
+		fileInfo: info,
+		swmr:     swmr,
+		etag:     info.ETag(),
 	}
+	sw := swmr.Writer()
 
 	go func() {
-		defer tee.Close()
-		defer resp.Close()
-		defer fw.Close()
-		defer swmr.Close()
+		defer body.Close()
+		_, err := io.Copy(sw, body)
+		_ = sw.CloseWithError(err)
+	}()
 
-		w := io.MultiWriter(swmr, fw)
-		n, err := io.Copy(w, body)
+	go func() {
+		r := swmr.NewReader(0)
+		defer r.Close()
+
+		defer fw.Close()
+		n, err := io.Copy(fw, r)
 		if err != nil && !errors.Is(err, io.EOF) {
 			if m.Logger != nil {
 				m.Logger.Println("SWMR copy error", cacheFile, contentLength, n, err)