feat(traffic): support limit task worker count & file stream rate (#7948)
* feat: set task workers num & client stream rate limit * feat: server stream rate limit * upgrade xhofe/tache * .
This commit is contained in:
KirCute_ECT
GitHub
152
internal/stream/limit.go
Normal file
152
internal/stream/limit.go
Normal file
@ -0,0 +1,152 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/alist-org/alist/v3/internal/model"
|
||||
"github.com/alist-org/alist/v3/pkg/http_range"
|
||||
"github.com/alist-org/alist/v3/pkg/utils"
|
||||
"golang.org/x/time/rate"
|
||||
"io"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Limiter interface {
|
||||
Limit() rate.Limit
|
||||
Burst() int
|
||||
TokensAt(time.Time) float64
|
||||
Tokens() float64
|
||||
Allow() bool
|
||||
AllowN(time.Time, int) bool
|
||||
Reserve() *rate.Reservation
|
||||
ReserveN(time.Time, int) *rate.Reservation
|
||||
Wait(context.Context) error
|
||||
WaitN(context.Context, int) error
|
||||
SetLimit(rate.Limit)
|
||||
SetLimitAt(time.Time, rate.Limit)
|
||||
SetBurst(int)
|
||||
SetBurstAt(time.Time, int)
|
||||
}
|
||||
|
||||
var (
|
||||
ClientDownloadLimit Limiter
|
||||
ClientUploadLimit Limiter
|
||||
ServerDownloadLimit Limiter
|
||||
ServerUploadLimit Limiter
|
||||
)
|
||||
|
||||
type RateLimitReader struct {
|
||||
io.Reader
|
||||
Limiter Limiter
|
||||
Ctx context.Context
|
||||
}
|
||||
|
||||
func (r *RateLimitReader) Read(p []byte) (n int, err error) {
|
||||
if r.Ctx != nil && utils.IsCanceled(r.Ctx) {
|
||||
return 0, r.Ctx.Err()
|
||||
}
|
||||
n, err = r.Reader.Read(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if r.Limiter != nil {
|
||||
if r.Ctx == nil {
|
||||
r.Ctx = context.Background()
|
||||
}
|
||||
err = r.Limiter.WaitN(r.Ctx, n)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (r *RateLimitReader) Close() error {
|
||||
if c, ok := r.Reader.(io.Closer); ok {
|
||||
return c.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type RateLimitWriter struct {
|
||||
io.Writer
|
||||
Limiter Limiter
|
||||
Ctx context.Context
|
||||
}
|
||||
|
||||
func (w *RateLimitWriter) Write(p []byte) (n int, err error) {
|
||||
if w.Ctx != nil && utils.IsCanceled(w.Ctx) {
|
||||
return 0, w.Ctx.Err()
|
||||
}
|
||||
n, err = w.Writer.Write(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if w.Limiter != nil {
|
||||
if w.Ctx == nil {
|
||||
w.Ctx = context.Background()
|
||||
}
|
||||
err = w.Limiter.WaitN(w.Ctx, n)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (w *RateLimitWriter) Close() error {
|
||||
if c, ok := w.Writer.(io.Closer); ok {
|
||||
return c.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type RateLimitFile struct {
|
||||
model.File
|
||||
Limiter Limiter
|
||||
Ctx context.Context
|
||||
}
|
||||
|
||||
func (r *RateLimitFile) Read(p []byte) (n int, err error) {
|
||||
if r.Ctx != nil && utils.IsCanceled(r.Ctx) {
|
||||
return 0, r.Ctx.Err()
|
||||
}
|
||||
n, err = r.File.Read(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if r.Limiter != nil {
|
||||
if r.Ctx == nil {
|
||||
r.Ctx = context.Background()
|
||||
}
|
||||
err = r.Limiter.WaitN(r.Ctx, n)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (r *RateLimitFile) ReadAt(p []byte, off int64) (n int, err error) {
|
||||
if r.Ctx != nil && utils.IsCanceled(r.Ctx) {
|
||||
return 0, r.Ctx.Err()
|
||||
}
|
||||
n, err = r.File.ReadAt(p, off)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if r.Limiter != nil {
|
||||
if r.Ctx == nil {
|
||||
r.Ctx = context.Background()
|
||||
}
|
||||
err = r.Limiter.WaitN(r.Ctx, n)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type RateLimitRangeReadCloser struct {
|
||||
model.RangeReadCloserIF
|
||||
Limiter Limiter
|
||||
}
|
||||
|
||||
func (rrc RateLimitRangeReadCloser) RangeRead(ctx context.Context, httpRange http_range.Range) (io.ReadCloser, error) {
|
||||
rc, err := rrc.RangeReadCloserIF.RangeRead(ctx, httpRange)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &RateLimitReader{
|
||||
Reader: rc,
|
||||
Limiter: rrc.Limiter,
|
||||
Ctx: ctx,
|
||||
}, nil
|
||||
}
|
||||
@ -182,14 +182,24 @@ func NewSeekableStream(fs FileStream, link *model.Link) (*SeekableStream, error)
|
||||
}
|
||||
if ss.Link != nil {
|
||||
if ss.Link.MFile != nil {
|
||||
ss.mFile = ss.Link.MFile
|
||||
ss.Reader = ss.Link.MFile
|
||||
ss.Closers.Add(ss.Link.MFile)
|
||||
mFile := ss.Link.MFile
|
||||
if _, ok := mFile.(*os.File); !ok {
|
||||
mFile = &RateLimitFile{
|
||||
File: mFile,
|
||||
Limiter: ServerDownloadLimit,
|
||||
Ctx: fs.Ctx,
|
||||
}
|
||||
}
|
||||
ss.mFile = mFile
|
||||
ss.Reader = mFile
|
||||
ss.Closers.Add(mFile)
|
||||
return &ss, nil
|
||||
}
|
||||
|
||||
if ss.Link.RangeReadCloser != nil {
|
||||
ss.rangeReadCloser = ss.Link.RangeReadCloser
|
||||
ss.rangeReadCloser = RateLimitRangeReadCloser{
|
||||
RangeReadCloserIF: ss.Link.RangeReadCloser,
|
||||
Limiter: ServerDownloadLimit,
|
||||
}
|
||||
ss.Add(ss.rangeReadCloser)
|
||||
return &ss, nil
|
||||
}
|
||||
@ -198,6 +208,10 @@ func NewSeekableStream(fs FileStream, link *model.Link) (*SeekableStream, error)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rrc = RateLimitRangeReadCloser{
|
||||
RangeReadCloserIF: rrc,
|
||||
Limiter: ServerDownloadLimit,
|
||||
}
|
||||
ss.rangeReadCloser = rrc
|
||||
ss.Add(rrc)
|
||||
return &ss, nil
|
||||
@ -259,7 +273,7 @@ func (ss *SeekableStream) CacheFullInTempFile() (model.File, error) {
|
||||
if ss.tmpFile != nil {
|
||||
return ss.tmpFile, nil
|
||||
}
|
||||
if ss.mFile != nil {
|
||||
if _, ok := ss.mFile.(*os.File); ok {
|
||||
return ss.mFile, nil
|
||||
}
|
||||
tmpF, err := utils.CreateTempFile(ss, ss.GetSize())
|
||||
@ -276,7 +290,7 @@ func (ss *SeekableStream) CacheFullInTempFileAndUpdateProgress(up model.UpdatePr
|
||||
if ss.tmpFile != nil {
|
||||
return ss.tmpFile, nil
|
||||
}
|
||||
if ss.mFile != nil {
|
||||
if _, ok := ss.mFile.(*os.File); ok {
|
||||
return ss.mFile, nil
|
||||
}
|
||||
tmpF, err := utils.CreateTempFile(&ReaderUpdatingProgress{
|
||||
@ -293,12 +307,13 @@ func (ss *SeekableStream) CacheFullInTempFileAndUpdateProgress(up model.UpdatePr
|
||||
}
|
||||
|
||||
func (f *FileStream) SetTmpFile(r *os.File) {
|
||||
f.Reader = r
|
||||
f.Add(r)
|
||||
f.tmpFile = r
|
||||
f.Reader = r
|
||||
}
|
||||
|
||||
type ReaderWithSize interface {
|
||||
io.Reader
|
||||
io.ReadCloser
|
||||
GetSize() int64
|
||||
}
|
||||
|
||||
@ -311,6 +326,13 @@ func (r *SimpleReaderWithSize) GetSize() int64 {
|
||||
return r.Size
|
||||
}
|
||||
|
||||
func (r *SimpleReaderWithSize) Close() error {
|
||||
if c, ok := r.Reader.(io.Closer); ok {
|
||||
return c.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type ReaderUpdatingProgress struct {
|
||||
Reader ReaderWithSize
|
||||
model.UpdateProgress
|
||||
@ -324,6 +346,10 @@ func (r *ReaderUpdatingProgress) Read(p []byte) (n int, err error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *ReaderUpdatingProgress) Close() error {
|
||||
return r.Reader.Close()
|
||||
}
|
||||
|
||||
type SStreamReadAtSeeker interface {
|
||||
model.File
|
||||
GetRawStream() *SeekableStream
|
||||
@ -534,7 +560,7 @@ func (r *RangeReadReadAtSeeker) Read(p []byte) (n int, err error) {
|
||||
|
||||
func (r *RangeReadReadAtSeeker) Close() error {
|
||||
if r.headCache != nil {
|
||||
r.headCache.close()
|
||||
_ = r.headCache.close()
|
||||
}
|
||||
return r.ss.Close()
|
||||
}
|
||||
@ -562,17 +588,3 @@ func (f *FileReadAtSeeker) Seek(offset int64, whence int) (int64, error) {
|
||||
func (f *FileReadAtSeeker) Close() error {
|
||||
return f.ss.Close()
|
||||
}
|
||||
|
||||
type ReaderWithCtx struct {
|
||||
io.Reader
|
||||
Ctx context.Context
|
||||
}
|
||||
|
||||
func (r *ReaderWithCtx) Read(p []byte) (n int, err error) {
|
||||
select {
|
||||
case <-r.Ctx.Done():
|
||||
return 0, r.Ctx.Err()
|
||||
default:
|
||||
return r.Reader.Read(p)
|
||||
}
|
||||
}
|
||||
|
||||
@ -3,6 +3,7 @@ package stream
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/alist-org/alist/v3/pkg/utils"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
@ -76,3 +77,22 @@ func checkContentRange(header *http.Header, offset int64) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type ReaderWithCtx struct {
|
||||
io.Reader
|
||||
Ctx context.Context
|
||||
}
|
||||
|
||||
func (r *ReaderWithCtx) Read(p []byte) (n int, err error) {
|
||||
if utils.IsCanceled(r.Ctx) {
|
||||
return 0, r.Ctx.Err()
|
||||
}
|
||||
return r.Reader.Read(p)
|
||||
}
|
||||
|
||||
func (r *ReaderWithCtx) Close() error {
|
||||
if c, ok := r.Reader.(io.Closer); ok {
|
||||
return c.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user