fix(stream): improve caching mechanism and handle size=0 case

Author: j2rong4cn

internal/stream/stream.go

@@ -4,18 +4,18 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/OpenListTeam/OpenList/v4/internal/conf"
-	"github.com/rclone/rclone/lib/mmap"
 	"io"
 	"math"
 	"os"
 	"sync"
 
+	"github.com/OpenListTeam/OpenList/v4/internal/conf"
 	"github.com/OpenListTeam/OpenList/v4/internal/errs"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
 	"github.com/OpenListTeam/OpenList/v4/pkg/buffer"
 	"github.com/OpenListTeam/OpenList/v4/pkg/http_range"
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
+	"github.com/rclone/rclone/lib/mmap"
 	"go4.org/readerutil"
 )
 
@@ -137,6 +137,61 @@ func (f *FileStream) CacheFullAndWriter(up *model.UpdateProgress, writer io.Writ
 	if writer != nil {
 		reader = io.TeeReader(reader, writer)
 	}
+
+	if f.GetSize() == 0 {
+		if f.peekBuff == nil {
+			f.peekBuff = &buffer.Reader{}
+		}
+		// 检查是否有数据
+		buf := make([]byte, 64*utils.KB)
+		n, err := io.ReadFull(reader, buf)
+		if err == io.ErrUnexpectedEOF {
+			if n > 0 {
+				f.peekBuff.Append(buf[:n])
+			}
+			f.size = f.peekBuff.Size()
+			f.Reader = f.peekBuff
+			return f.peekBuff, nil
+		} else if err != nil {
+			return nil, err
+		}
+		f.peekBuff.Append(buf[:n])
+		if conf.MaxBufferLimit-n > conf.MmapThreshold && conf.MmapThreshold > 0 {
+			m, err := mmap.Alloc(conf.MaxBufferLimit - n)
+			if err == nil {
+				f.Add(utils.CloseFunc(func() error {
+					return mmap.Free(m)
+				}))
+				n, err = io.ReadFull(reader, m)
+				if err == io.ErrUnexpectedEOF {
+					if n > 0 {
+						f.peekBuff.Append(m[:n])
+					}
+					f.size = f.peekBuff.Size()
+					f.Reader = f.peekBuff
+					return f.peekBuff, nil
+				} else if err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		tmpF, err := utils.CreateTempFile(reader, 0)
+		if err != nil {
+			return nil, err
+		}
+		f.Add(utils.CloseFunc(func() error {
+			return errors.Join(tmpF.Close(), os.RemoveAll(tmpF.Name()))
+		}))
+		peekF, err := buffer.NewPeekFile(f.peekBuff, tmpF)
+		if err != nil {
+			return nil, err
+		}
+		f.size = peekF.Size()
+		f.Reader = peekF
+		return peekF, nil
+	}
+
 	f.Reader = reader
 	return f.cache(f.GetSize())
 }
@@ -162,7 +217,7 @@ func (f *FileStream) RangeRead(httpRange http_range.Range) (io.Reader, error) {
 	}
 
 	size := httpRange.Start + httpRange.Length
-	if f.peekBuff != nil && size <= int64(f.peekBuff.Len()) {
+	if f.peekBuff != nil && size <= int64(f.peekBuff.Size()) {
 		return io.NewSectionReader(f.peekBuff, httpRange.Start, httpRange.Length), nil
 	}
 
@@ -179,57 +234,47 @@ func (f *FileStream) RangeRead(httpRange http_range.Range) (io.Reader, error) {
 // 即使被写入的数据量与Buffer.Cap一致，Buffer也会扩大
 
 func (f *FileStream) cache(maxCacheSize int64) (model.File, error) {
-	limit := int64(conf.MaxBufferLimit)
+	if maxCacheSize > int64(conf.MaxBufferLimit) {
+		tmpF, err := utils.CreateTempFile(f.Reader, f.GetSize())
+		if err != nil {
+			return nil, err
+		}
+		f.Add(tmpF)
+		f.tmpFile = tmpF
+		f.Reader = tmpF
+		return tmpF, nil
+	}
 
 	if f.peekBuff == nil {
 		f.peekBuff = &buffer.Reader{}
 		f.oriReader = f.Reader
 	}
+	bufSize := maxCacheSize - int64(f.peekBuff.Size())
 	var buf []byte
-	bufSize := 64 << 10 // 64KB as default
-	if conf.MmapThreshold > 0 && bufSize >= conf.MmapThreshold {
-		m, err := mmap.Alloc(bufSize)
+	if conf.MmapThreshold > 0 && bufSize >= int64(conf.MmapThreshold) {
+		m, err := mmap.Alloc(int(bufSize))
 		if err == nil {
 			f.Add(utils.CloseFunc(func() error {
 				return mmap.Free(m)
 			}))
 			buf = m
 		}
 	}
-
-	var readBytes int
-	// precache first `limit` byte
-	for int64(readBytes) < limit {
-		if buf == nil {
-			buf = make([]byte, bufSize)
-		}
-
-		want := limit - int64(readBytes)
-		if want > int64(len(buf)) {
-			want = int64(len(buf))
-		}
-		n, err := f.oriReader.Read(buf[:want])
-		if n > 0 {
-			f.peekBuff.Append(buf[:n])
-			readBytes += n
-		}
-		if err == io.EOF {
-			f.Reader = f.peekBuff
-			f.oriReader = nil
-			// should update real file size here to solve `GetSize == 0` issue
-			f.size = int64(readBytes)
-			return f.peekBuff, nil
-		}
+	if buf == nil {
+		buf = make([]byte, bufSize)
 	}
-	// if file is larger than MaxBufferLimit, fallback to disk
-	tmpF, err := utils.CreateTempFile(io.MultiReader(f.peekBuff, f.oriReader), f.GetSize())
-	if err != nil {
-		return nil, err
+	n, err := io.ReadFull(f.oriReader, buf)
+	if bufSize != int64(n) {
+		return nil, fmt.Errorf("failed to read all data: (expect =%d, actual =%d) %w", bufSize, n, err)
+	}
+	f.peekBuff.Append(buf)
+	if int64(f.peekBuff.Size()) >= f.GetSize() {
+		f.Reader = f.peekBuff
+		f.oriReader = nil
+	} else {
+		f.Reader = io.MultiReader(f.peekBuff, f.oriReader)
 	}
-	f.Add(tmpF)
-	f.tmpFile = tmpF
-	f.Reader = tmpF
-	return tmpF, nil
+	return f.peekBuff, nil
 }
 
 func (f *FileStream) SetTmpFile(file model.File) {

pkg/buffer/bytes.go

@@ -8,29 +8,29 @@ import (
 // 用于存储不复用的[]byte
 type Reader struct {
 	bufs   [][]byte
-	length int
-	offset int
+	size   int64
+	offset int64
 }
 
-func (r *Reader) Len() int {
-	return r.length
+func (r *Reader) Size() int64 {
+	return r.size
 }
 
 func (r *Reader) Append(buf []byte) {
-	r.length += len(buf)
+	r.size += int64(len(buf))
 	r.bufs = append(r.bufs, buf)
 }
 
 func (r *Reader) Read(p []byte) (int, error) {
-	n, err := r.ReadAt(p, int64(r.offset))
+	n, err := r.ReadAt(p, r.offset)
 	if n > 0 {
-		r.offset += n
+		r.offset += int64(n)
 	}
 	return n, err
 }
 
 func (r *Reader) ReadAt(p []byte, off int64) (int, error) {
-	if off < 0 || off >= int64(r.length) {
+	if off < 0 || off >= r.size {
 		return 0, io.EOF
 	}
 
@@ -56,35 +56,35 @@ func (r *Reader) ReadAt(p []byte, off int64) (int, error) {
 }
 
 func (r *Reader) Seek(offset int64, whence int) (int64, error) {
-	var abs int
 	switch whence {
 	case io.SeekStart:
-		abs = int(offset)
 	case io.SeekCurrent:
-		abs = r.offset + int(offset)
+		offset = r.offset + offset
 	case io.SeekEnd:
-		abs = r.length + int(offset)
+		offset = r.size + offset
 	default:
 		return 0, errors.New("Seek: invalid whence")
 	}
 
-	if abs < 0 || abs > r.length {
+	if offset < 0 || offset > r.size {
 		return 0, errors.New("Seek: invalid offset")
 	}
 
-	r.offset = abs
-	return int64(abs), nil
+	r.offset = offset
+	return offset, nil
 }
 
 func (r *Reader) Reset() {
 	clear(r.bufs)
 	r.bufs = nil
-	r.length = 0
+	r.size = 0
 	r.offset = 0
 }
 
 func NewReader(buf ...[]byte) *Reader {
-	b := &Reader{}
+	b := &Reader{
+		bufs: make([][]byte, 0, len(buf)),
+	}
 	for _, b1 := range buf {
 		b.Append(b1)
 	}

pkg/buffer/bytes_test.go

@@ -71,7 +71,7 @@ func TestReader_ReadAt(t *testing.T) {
 				off: 24,
 			},
 			want: func(a args, n int, err error) error {
-				if n != bs.Len()-int(a.off) {
+				if n != int(bs.Size()-a.off) {
 					return errors.New("read length not match")
 				}
 				if string(a.p[:n]) != "OpenList" {

pkg/buffer/file.go

@@ -0,0 +1,88 @@
+package buffer
+
+import (
+	"errors"
+	"io"
+	"os"
+)
+
+type PeekFile struct {
+	peek   *Reader
+	file   *os.File
+	offset int64
+	size   int64
+}
+
+func (p *PeekFile) Read(b []byte) (n int, err error) {
+	n, err = p.ReadAt(b, p.offset)
+	if n > 0 {
+		p.offset += int64(n)
+	}
+	return n, err
+}
+
+func (p *PeekFile) ReadAt(b []byte, off int64) (n int, err error) {
+	if off < p.peek.Size() {
+		n, err = p.peek.ReadAt(b, off)
+		if err == nil || n == len(b) {
+			return n, nil
+		}
+		// EOF
+	}
+	var nn int
+	nn, err = p.file.ReadAt(b[n:], off+int64(n)-p.peek.Size())
+	return n + nn, err
+}
+
+func (p *PeekFile) Seek(offset int64, whence int) (int64, error) {
+	switch whence {
+	case io.SeekStart:
+	case io.SeekCurrent:
+		if offset == 0 {
+			return p.offset, nil
+		}
+		offset = p.offset + offset
+	case io.SeekEnd:
+		offset = p.size + offset
+	default:
+		return 0, errors.New("Seek: invalid whence")
+	}
+
+	if offset < 0 || offset > p.size {
+		return 0, errors.New("Seek: invalid offset")
+	}
+	if offset <= p.peek.Size() {
+		_, err := p.peek.Seek(offset, io.SeekStart)
+		if err != nil {
+			return 0, err
+		}
+		_, err = p.file.Seek(0, io.SeekStart)
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		_, err := p.peek.Seek(p.peek.Size(), io.SeekStart)
+		if err != nil {
+			return 0, err
+		}
+		_, err = p.file.Seek(offset-p.peek.Size(), io.SeekStart)
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	p.offset = offset
+	return offset, nil
+}
+
+func (p *PeekFile) Size() int64 {
+	return p.size
+}
+
+func NewPeekFile(peek *Reader, file *os.File) (*PeekFile, error) {
+	stat, err := file.Stat()
+	if err == nil {
+		return &PeekFile{peek: peek, file: file, size: stat.Size() + peek.Size()}, nil
+	}
+	return nil, err
+}