ProjectHentai/alist
6
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(traffic): support limit task worker count & file stream rate (#7948)

Browse Source 
* feat: set task workers num & client stream rate limit

* feat: server stream rate limit

* upgrade xhofe/tache

* .
This commit is contained in:
KirCute_ECT
2025-02-16 12:22:11 +08:00
committed by GitHub
parent 399336b33c
commit 3b71500f23
79 changed files with 803 additions and 327 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
server/common/proxy.go
View File

@ -6,6 +6,7 @@ import (
"io"
"net/http"
"net/url"
"os"
"github.com/alist-org/alist/v3/internal/model"
"github.com/alist-org/alist/v3/internal/net"
@ -23,11 +24,22 @@ func Proxy(w http.ResponseWriter, r *http.Request, link *model.Link, file model.
if contentType != "" {
w.Header().Set("Content-Type", contentType)
}
http.ServeContent(w, r, file.GetName(), file.ModTime(), link.MFile)
mFile := link.MFile
if _, ok := mFile.(*os.File); !ok {
mFile = &stream.RateLimitFile{
File: mFile,
Limiter: stream.ServerDownloadLimit,
Ctx: r.Context(),
}
}
http.ServeContent(w, r, file.GetName(), file.ModTime(), mFile)
return nil
} else if link.RangeReadCloser != nil {
attachFileName(w, file)
net.ServeHTTP(w, r, file.GetName(), file.ModTime(), file.GetSize(), link.RangeReadCloser)
net.ServeHTTP(w, r, file.GetName(), file.ModTime(), file.GetSize(), &stream.RateLimitRangeReadCloser{
RangeReadCloserIF: link.RangeReadCloser,
Limiter: stream.ServerDownloadLimit,
})
return nil
} else if link.Concurrency != 0 || link.PartSize != 0 {
attachFileName(w, file)
@ -47,7 +59,10 @@ func Proxy(w http.ResponseWriter, r *http.Request, link *model.Link, file model.
rc, err := down.Download(ctx, req)
return rc, err
}
net.ServeHTTP(w, r, file.GetName(), file.ModTime(), file.GetSize(), &model.RangeReadCloser{RangeReader: rangeReader})
net.ServeHTTP(w, r, file.GetName(), file.ModTime(), file.GetSize(), &stream.RateLimitRangeReadCloser{
RangeReadCloserIF: &model.RangeReadCloser{RangeReader: rangeReader},
Limiter: stream.ServerDownloadLimit,
})
return nil
} else {
//transparent proxy
@ -65,7 +80,11 @@ func Proxy(w http.ResponseWriter, r *http.Request, link *model.Link, file model.
if r.Method == http.MethodHead {
return nil
}
_, err = utils.CopyWithBuffer(w, res.Body)
_, err = utils.CopyWithBuffer(w, &stream.RateLimitReader{
Reader: res.Body,
Limiter: stream.ServerDownloadLimit,
Ctx: r.Context(),
})
if err != nil {
return err
}

7
server/ftp/fsread.go
View File

@ -60,7 +60,12 @@ func OpenDownload(ctx context.Context, reqPath string, offset int64) (*FileDownl
}
func (f *FileDownloadProxy) Read(p []byte) (n int, err error) {
return f.reader.Read(p)
n, err = f.reader.Read(p)
if err != nil {
return
}
err = stream.ClientDownloadLimit.WaitN(f.reader.GetRawStream().Ctx, n)
return
}
func (f *FileDownloadProxy) Write(p []byte) (n int, err error) {

19
server/ftp/fsup.go
View File

@ -59,7 +59,12 @@ func (f *FileUploadProxy) Read(p []byte) (n int, err error) {
}
func (f *FileUploadProxy) Write(p []byte) (n int, err error) {
return f.buffer.Write(p)
n, err = f.buffer.Write(p)
if err != nil {
return
}
err = stream.ClientUploadLimit.WaitN(f.ctx, n)
return
}
func (f *FileUploadProxy) Seek(offset int64, whence int) (int64, error) {
@ -96,7 +101,6 @@ func (f *FileUploadProxy) Close() error {
WebPutAsTask: true,
}
s.SetTmpFile(f.buffer)
s.Closers.Add(f.buffer)
_, err = fs.PutAsTask(f.ctx, dir, s)
return err
}
@ -127,7 +131,7 @@ func (f *FileUploadWithLengthProxy) Read(p []byte) (n int, err error) {
return 0, errs.NotSupport
}
func (f *FileUploadWithLengthProxy) Write(p []byte) (n int, err error) {
func (f *FileUploadWithLengthProxy) write(p []byte) (n int, err error) {
if f.pipeWriter != nil {
select {
case e := <-f.errChan:
@ -174,6 +178,15 @@ func (f *FileUploadWithLengthProxy) Write(p []byte) (n int, err error) {
}
}
func (f *FileUploadWithLengthProxy) Write(p []byte) (n int, err error) {
n, err = f.write(p)
if err != nil {
return
}
err = stream.ClientUploadLimit.WaitN(f.ctx, n)
return
}
func (f *FileUploadWithLengthProxy) Seek(offset int64, whence int) (int64, error) {
return 0, errs.NotSupport
}

36
server/middlewares/limit.go
View File

@ -1,7 +1,9 @@
package middlewares
import (
"github.com/alist-org/alist/v3/internal/stream"
"github.com/gin-gonic/gin"
"io"
)
func MaxAllowed(n int) gin.HandlerFunc {
@ -14,3 +16,37 @@ func MaxAllowed(n int) gin.HandlerFunc {
c.Next()
}
}
func UploadRateLimiter(limiter stream.Limiter) gin.HandlerFunc {
return func(c *gin.Context) {
c.Request.Body = &stream.RateLimitReader{
Reader: c.Request.Body,
Limiter: limiter,
Ctx: c,
}
c.Next()
}
}
type ResponseWriterWrapper struct {
gin.ResponseWriter
WrapWriter io.Writer
}
func (w *ResponseWriterWrapper) Write(p []byte) (n int, err error) {
return w.WrapWriter.Write(p)
}
func DownloadRateLimiter(limiter stream.Limiter) gin.HandlerFunc {
return func(c *gin.Context) {
c.Writer = &ResponseWriterWrapper{
ResponseWriter: c.Writer,
WrapWriter: &stream.RateLimitWriter{
Writer: c.Writer,
Limiter: limiter,
Ctx: c,
},
}
c.Next()
}
}

17
server/router.go
View File

@ -4,6 +4,7 @@ import (
"github.com/alist-org/alist/v3/cmd/flags"
"github.com/alist-org/alist/v3/internal/conf"
"github.com/alist-org/alist/v3/internal/message"
"github.com/alist-org/alist/v3/internal/stream"
"github.com/alist-org/alist/v3/pkg/utils"
"github.com/alist-org/alist/v3/server/common"
"github.com/alist-org/alist/v3/server/handles"
@ -38,13 +39,14 @@ func Init(e *gin.Engine) {
WebDav(g.Group("/dav"))
S3(g.Group("/s3"))
g.GET("/d/*path", middlewares.Down, handles.Down)
g.GET("/p/*path", middlewares.Down, handles.Proxy)
downloadLimiter := middlewares.DownloadRateLimiter(stream.ClientDownloadLimit)
g.GET("/d/*path", middlewares.Down, downloadLimiter, handles.Down)
g.GET("/p/*path", middlewares.Down, downloadLimiter, handles.Proxy)
g.HEAD("/d/*path", middlewares.Down, handles.Down)
g.HEAD("/p/*path", middlewares.Down, handles.Proxy)
g.GET("/ad/*path", middlewares.Down, handles.ArchiveDown)
g.GET("/ap/*path", middlewares.Down, handles.ArchiveProxy)
g.GET("/ae/*path", middlewares.Down, handles.ArchiveInternalExtract)
g.GET("/ad/*path", middlewares.Down, downloadLimiter, handles.ArchiveDown)
g.GET("/ap/*path", middlewares.Down, downloadLimiter, handles.ArchiveProxy)
g.GET("/ae/*path", middlewares.Down, downloadLimiter, handles.ArchiveInternalExtract)
g.HEAD("/ad/*path", middlewares.Down, handles.ArchiveDown)
g.HEAD("/ap/*path", middlewares.Down, handles.ArchiveProxy)
g.HEAD("/ae/*path", middlewares.Down, handles.ArchiveInternalExtract)
@ -173,8 +175,9 @@ func _fs(g *gin.RouterGroup) {
g.POST("/copy", handles.FsCopy)
g.POST("/remove", handles.FsRemove)
g.POST("/remove_empty_directory", handles.FsRemoveEmptyDirectory)
g.PUT("/put", middlewares.FsUp, handles.FsStream)
g.PUT("/form", middlewares.FsUp, handles.FsForm)
uploadLimiter := middlewares.UploadRateLimiter(stream.ClientUploadLimit)
g.PUT("/put", middlewares.FsUp, uploadLimiter, handles.FsStream)
g.PUT("/form", middlewares.FsUp, uploadLimiter, handles.FsForm)
g.POST("/link", middlewares.AuthAdmin, handles.Link)
// g.POST("/add_aria2", handles.AddOfflineDownload)
// g.POST("/add_qbit", handles.AddQbittorrent)

8
server/webdav.go
View File

@ -3,6 +3,8 @@ package server
import (
"context"
"crypto/subtle"
"github.com/alist-org/alist/v3/internal/stream"
"github.com/alist-org/alist/v3/server/middlewares"
"net/http"
"path"
"strings"
@ -27,8 +29,10 @@ func WebDav(dav *gin.RouterGroup) {
},
}
dav.Use(WebDAVAuth)
dav.Any("/*path", ServeWebDAV)
dav.Any("", ServeWebDAV)
uploadLimiter := middlewares.UploadRateLimiter(stream.ClientUploadLimit)
downloadLimiter := middlewares.DownloadRateLimiter(stream.ClientDownloadLimit)
dav.Any("/*path", uploadLimiter, downloadLimiter, ServeWebDAV)
dav.Any("", uploadLimiter, downloadLimiter, ServeWebDAV)
dav.Handle("PROPFIND", "/*path", ServeWebDAV)
dav.Handle("PROPFIND", "", ServeWebDAV)
dav.Handle("MKCOL", "/*path", ServeWebDAV)