feat(stream): fast buffer freeing for large cache (#1053)

Signed-off-by: j2rong4cn <36783515+j2rong4cn@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

internal/bootstrap/config.go

@@ -77,6 +77,10 @@ func InitConfig() {
 			log.Fatalf("update config struct error: %+v", err)
 		}
 	}
+	if !conf.Conf.Force {
+		confFromEnv()
+	}
+
 	if conf.Conf.MaxConcurrency > 0 {
 		net.DefaultConcurrencyLimit = &net.ConcurrencyLimit{Limit: conf.Conf.MaxConcurrency}
 	}
@@ -92,25 +96,31 @@ func InitConfig() {
 		conf.MaxBufferLimit = conf.Conf.MaxBufferLimit * utils.MB
 	}
 	log.Infof("max buffer limit: %dMB", conf.MaxBufferLimit/utils.MB)
-	if !conf.Conf.Force {
+	if conf.Conf.MmapThreshold > 0 {
-		confFromEnv()
+		conf.MmapThreshold = conf.Conf.MmapThreshold * utils.MB
+	} else {
+		conf.MmapThreshold = 0
 	}
+	log.Infof("mmap threshold: %dMB", conf.Conf.MmapThreshold)
+
 	if len(conf.Conf.Log.Filter.Filters) == 0 {
 		conf.Conf.Log.Filter.Enable = false
 	}
 	// convert abs path
 	convertAbsPath := func(path *string) {
-		if !filepath.IsAbs(*path) {
+		if *path != "" && !filepath.IsAbs(*path) {
 			*path = filepath.Join(pwd, *path)
 		}
 	}
+	convertAbsPath(&conf.Conf.Database.DBFile)
+	convertAbsPath(&conf.Conf.Scheme.CertFile)
+	convertAbsPath(&conf.Conf.Scheme.KeyFile)
+	convertAbsPath(&conf.Conf.Scheme.UnixFile)
+	convertAbsPath(&conf.Conf.Log.Name)
 	convertAbsPath(&conf.Conf.TempDir)
 	convertAbsPath(&conf.Conf.BleveDir)
-	convertAbsPath(&conf.Conf.Log.Name)
-	convertAbsPath(&conf.Conf.Database.DBFile)
-	if conf.Conf.DistDir != "" {
 	convertAbsPath(&conf.Conf.DistDir)
-	}
+
 	err := os.MkdirAll(conf.Conf.TempDir, 0o777)
 	if err != nil {
 		log.Fatalf("create temp dir error: %+v", err)

internal/conf/config.go

@@ -120,6 +120,7 @@ type Config struct {
 	Log                   LogConfig   `json:"log" envPrefix:"LOG_"`
 	DelayedStart          int         `json:"delayed_start" env:"DELAYED_START"`
 	MaxBufferLimit        int         `json:"max_buffer_limitMB" env:"MAX_BUFFER_LIMIT_MB"`
+	MmapThreshold         int         `json:"mmap_thresholdMB" env:"MMAP_THRESHOLD_MB"`
 	MaxConnections        int         `json:"max_connections" env:"MAX_CONNECTIONS"`
 	MaxConcurrency        int         `json:"max_concurrency" env:"MAX_CONCURRENCY"`
 	TlsInsecureSkipVerify bool        `json:"tls_insecure_skip_verify" env:"TLS_INSECURE_SKIP_VERIFY"`
@@ -176,6 +177,7 @@ func DefaultConfig(dataDir string) *Config {
 			},
 		},
 		MaxBufferLimit:        -1,
+		MmapThreshold:         4,
 		MaxConnections:        0,
 		MaxConcurrency:        64,
 		TlsInsecureSkipVerify: true,

internal/conf/var.go

@@ -25,7 +25,10 @@ var PrivacyReg []*regexp.Regexp
 var (
 	// StoragesLoaded loaded success if empty
 	StoragesLoaded = false
+	// 单个Buffer最大限制
 	MaxBufferLimit = 16 * 1024 * 1024
+	// 超过该阈值的Buffer将使用 mmap 分配，可主动释放内存
+	MmapThreshold = 4 * 1024 * 1024
 )
 var (
 	RawIndexHtml string

internal/net/request.go

@@ -1,7 +1,6 @@
 package net
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -15,6 +14,7 @@ import (
 	"github.com/OpenListTeam/OpenList/v4/internal/conf"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
+	"github.com/rclone/rclone/lib/mmap"
 
 	"github.com/OpenListTeam/OpenList/v4/pkg/http_range"
 	"github.com/aws/aws-sdk-go/aws/awsutil"
@@ -255,7 +255,10 @@ func (d *downloader) sendChunkTask(newConcurrency bool) error {
 				finalSize += firstSize - minSize
 			}
 		}
-		buf.Reset(int(finalSize))
+		err := buf.Reset(int(finalSize))
+		if err != nil {
+			return err
+		}
 		ch := chunk{
 			start: d.pos,
 			size:  finalSize,
@@ -645,11 +648,13 @@ func (mr MultiReadCloser) Close() error {
 }
 
 type Buf struct {
-	buffer *bytes.Buffer
 	size int //expected size
 	ctx  context.Context
-	off    int
+	offR int
+	offW int
 	rw   sync.Mutex
+	buf  []byte
+	mmap bool
 
 	readSignal  chan struct{}
 	readPending bool
@@ -658,54 +663,62 @@ type Buf struct {
 // NewBuf is a buffer that can have 1 read & 1 write at the same time.
 // when read is faster write, immediately feed data to read after written
 func NewBuf(ctx context.Context, maxSize int) *Buf {
-	return &Buf{
+	br := &Buf{
 		ctx:        ctx,
-		buffer: bytes.NewBuffer(make([]byte, 0, maxSize)),
 		size:       maxSize,
-
 		readSignal: make(chan struct{}, 1),
 	}
+	if conf.MmapThreshold > 0 && maxSize >= conf.MmapThreshold {
+		m, err := mmap.Alloc(maxSize)
+		if err == nil {
+			br.buf = m
+			br.mmap = true
+			return br
 		}
-func (br *Buf) Reset(size int) {
-	br.rw.Lock()
-	defer br.rw.Unlock()
-	if br.buffer == nil {
-		return
 	}
-	br.buffer.Reset()
+	br.buf = make([]byte, maxSize)
-	br.size = size
+	return br
-	br.off = 0
 }
 
-func (br *Buf) Read(p []byte) (n int, err error) {
+func (br *Buf) Reset(size int) error {
+	br.rw.Lock()
+	defer br.rw.Unlock()
+	if br.buf == nil {
+		return io.ErrClosedPipe
+	}
+	if size > cap(br.buf) {
+		return fmt.Errorf("reset size %d exceeds max size %d", size, cap(br.buf))
+	}
+	br.size = size
+	br.offR = 0
+	br.offW = 0
+	return nil
+}
+
+func (br *Buf) Read(p []byte) (int, error) {
 	if err := br.ctx.Err(); err != nil {
 		return 0, err
 	}
 	if len(p) == 0 {
 		return 0, nil
 	}
-	if br.off >= br.size {
+	if br.offR >= br.size {
 		return 0, io.EOF
 	}
 	for {
 		br.rw.Lock()
-		if br.buffer != nil {
+		if br.buf == nil {
-			n, err = br.buffer.Read(p)
-		} else {
-			err = io.ErrClosedPipe
-		}
-		if err != nil && err != io.EOF {
 			br.rw.Unlock()
-			return
+			return 0, io.ErrClosedPipe
 		}
-		if n > 0 {
+
-			br.off += n
+		if br.offW < br.offR {
 			br.rw.Unlock()
-			return n, nil
+			return 0, io.ErrUnexpectedEOF
 		}
+		if br.offW == br.offR {
 			br.readPending = true
 			br.rw.Unlock()
-		// n==0, err==io.EOF
 			select {
 			case <-br.ctx.Done():
 				return 0, br.ctx.Err()
@@ -716,18 +729,34 @@ func (br *Buf) Read(p []byte) (n int, err error) {
 				continue
 			}
 		}
+
+		n := copy(p, br.buf[br.offR:br.offW])
+		br.offR += n
+		br.rw.Unlock()
+		if n < len(p) && br.offR >= br.size {
+			return n, io.EOF
+		}
+		return n, nil
+	}
 }
 
-func (br *Buf) Write(p []byte) (n int, err error) {
+func (br *Buf) Write(p []byte) (int, error) {
 	if err := br.ctx.Err(); err != nil {
 		return 0, err
 	}
+	if len(p) == 0 {
+		return 0, nil
+	}
 	br.rw.Lock()
 	defer br.rw.Unlock()
-	if br.buffer == nil {
+	if br.buf == nil {
 		return 0, io.ErrClosedPipe
 	}
-	n, err = br.buffer.Write(p)
+	if br.offW >= br.size {
+		return 0, io.ErrShortWrite
+	}
+	n := copy(br.buf[br.offW:], p[:min(br.size-br.offW, len(p))])
+	br.offW += n
 	if br.readPending {
 		br.readPending = false
 		select {
@@ -735,12 +764,21 @@ func (br *Buf) Write(p []byte) (n int, err error) {
 		default:
 		}
 	}
-	return
+	if n < len(p) {
+		return n, io.ErrShortWrite
+	}
+	return n, nil
 }
 
-func (br *Buf) Close() {
+func (br *Buf) Close() error {
 	br.rw.Lock()
 	defer br.rw.Unlock()
-	br.buffer = nil
+	var err error
-	close(br.readSignal)
+	if br.mmap {
+		err = mmap.Free(br.buf)
+		br.mmap = false
+	}
+	br.buf = nil
+	close(br.readSignal)
+	return err
 }

internal/stream/stream.go

@@ -15,6 +15,7 @@ import (
 	"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"
 )
 
@@ -60,8 +61,12 @@ func (f *FileStream) IsForceStreamUpload() bool {
 }
 
 func (f *FileStream) Close() error {
-	var err1, err2 error
+	if f.peekBuff != nil {
+		f.peekBuff.Reset()
+		f.peekBuff = nil
+	}
 
+	var err1, err2 error
 	err1 = f.Closers.Close()
 	if errors.Is(err1, os.ErrClosed) {
 		err1 = nil
@@ -74,10 +79,6 @@ func (f *FileStream) Close() error {
 			f.tmpFile = nil
 		}
 	}
-	if f.peekBuff != nil {
-		f.peekBuff.Reset()
-		f.peekBuff = nil
-	}
 
 	return errors.Join(err1, err2)
 }
@@ -194,7 +195,19 @@ func (f *FileStream) cache(maxCacheSize int64) (model.File, error) {
 		f.oriReader = f.Reader
 	}
 	bufSize := maxCacheSize - int64(f.peekBuff.Len())
-	buf := make([]byte, bufSize)
+	var buf []byte
+	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
+		}
+	}
+	if buf == nil {
+		buf = make([]byte, bufSize)
+	}
 	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)

internal/stream/stream_test.go

@@ -7,11 +7,13 @@ import (
 	"io"
 	"testing"
 
+	"github.com/OpenListTeam/OpenList/v4/internal/conf"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
 	"github.com/OpenListTeam/OpenList/v4/pkg/http_range"
 )
 
 func TestFileStream_RangeRead(t *testing.T) {
+	conf.MaxBufferLimit = 16 * 1024 * 1024
 	type args struct {
 		httpRange http_range.Range
 	}
@@ -71,7 +73,7 @@ func TestFileStream_RangeRead(t *testing.T) {
 			}
 		})
 	}
-	t.Run("after check", func(t *testing.T) {
+	t.Run("after", func(t *testing.T) {
 		if f.GetFile() == nil {
 			t.Error("not cached")
 		}

internal/stream/util.go

@@ -8,13 +8,14 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-	"sync"
 
 	"github.com/OpenListTeam/OpenList/v4/internal/conf"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
 	"github.com/OpenListTeam/OpenList/v4/internal/net"
 	"github.com/OpenListTeam/OpenList/v4/pkg/http_range"
+	"github.com/OpenListTeam/OpenList/v4/pkg/pool"
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
+	"github.com/rclone/rclone/lib/mmap"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -153,26 +154,49 @@ func CacheFullAndHash(stream model.FileStreamer, up *model.UpdateProgress, hashT
 type StreamSectionReader struct {
 	file    model.FileStreamer
 	off     int64
-	bufPool *sync.Pool
+	bufPool *pool.Pool[[]byte]
 }
 
 func NewStreamSectionReader(file model.FileStreamer, maxBufferSize int, up *model.UpdateProgress) (*StreamSectionReader, error) {
 	ss := &StreamSectionReader{file: file}
-	if file.GetFile() == nil {
+	if file.GetFile() != nil {
+		return ss, nil
+	}
+
 	maxBufferSize = min(maxBufferSize, int(file.GetSize()))
 	if maxBufferSize > conf.MaxBufferLimit {
 		_, err := file.CacheFullAndWriter(up, nil)
 		if err != nil {
 			return nil, err
 		}
+		return ss, nil
+	}
+	if conf.MmapThreshold > 0 && maxBufferSize >= conf.MmapThreshold {
+		ss.bufPool = &pool.Pool[[]byte]{
+			New: func() []byte {
+				buf, err := mmap.Alloc(maxBufferSize)
+				if err == nil {
+					file.Add(utils.CloseFunc(func() error {
+						return mmap.Free(buf)
+					}))
 				} else {
-			ss.bufPool = &sync.Pool{
+					buf = make([]byte, maxBufferSize)
-				New: func() any {
+				}
+				return buf
+			},
+		}
+	} else {
+		ss.bufPool = &pool.Pool[[]byte]{
+			New: func() []byte {
 				return make([]byte, maxBufferSize)
 			},
 		}
 	}
-	}
+
+	file.Add(utils.CloseFunc(func() error {
+		ss.bufPool.Reset()
+		return nil
+	}))
 	return ss, nil
 }
 
@@ -184,7 +208,7 @@ func (ss *StreamSectionReader) GetSectionReader(off, length int64) (*SectionRead
 		if off != ss.off {
 			return nil, fmt.Errorf("stream not cached: request offset %d != current offset %d", off, ss.off)
 		}
-		tempBuf := ss.bufPool.Get().([]byte)
+		tempBuf := ss.bufPool.Get()
 		buf = tempBuf[:length]
 		n, err := io.ReadFull(ss.file, buf)
 		if int64(n) != length {

pkg/pool/pool.go

@@ -0,0 +1,37 @@
+package pool
+
+import "sync"
+
+type Pool[T any] struct {
+	New    func() T
+	MaxCap int
+
+	cache []T
+	mu    sync.Mutex
+}
+
+func (p *Pool[T]) Get() T {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if len(p.cache) == 0 {
+		return p.New()
+	}
+	item := p.cache[len(p.cache)-1]
+	p.cache = p.cache[:len(p.cache)-1]
+	return item
+}
+
+func (p *Pool[T]) Put(item T) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.MaxCap == 0 || len(p.cache) < int(p.MaxCap) {
+		p.cache = append(p.cache, item)
+	}
+}
+
+func (p *Pool[T]) Reset() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	clear(p.cache)
+	p.cache = nil
+}