mirror of
https://github.com/TeaOSLab/EdgeNode.git
synced 2025-11-14 15:40:24 +08:00
优化请求限制逻辑,连接关闭时自动终止内容发送
This commit is contained in:
@@ -132,7 +132,7 @@ func (this *HTTPWriter) Prepare(resp *http.Response, size int64, status int, ena
|
|||||||
this.req.web.RequestLimit != nil &&
|
this.req.web.RequestLimit != nil &&
|
||||||
this.req.web.RequestLimit.IsOn &&
|
this.req.web.RequestLimit.IsOn &&
|
||||||
this.req.web.RequestLimit.OutBandwidthPerConnBytes() > 0 {
|
this.req.web.RequestLimit.OutBandwidthPerConnBytes() > 0 {
|
||||||
this.writer = writers.NewRateLimitWriter(this.writer, this.req.web.RequestLimit.OutBandwidthPerConnBytes())
|
this.writer = writers.NewRateLimitWriter(this.req.RawReq.Context(), this.writer, this.req.web.RequestLimit.OutBandwidthPerConnBytes())
|
||||||
}
|
}
|
||||||
|
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -3,6 +3,7 @@
|
|||||||
package writers
|
package writers
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"github.com/iwind/TeaGo/types"
|
"github.com/iwind/TeaGo/types"
|
||||||
"io"
|
"io"
|
||||||
"time"
|
"time"
|
||||||
@@ -11,6 +12,7 @@ import (
|
|||||||
// RateLimitWriter 限速写入
|
// RateLimitWriter 限速写入
|
||||||
type RateLimitWriter struct {
|
type RateLimitWriter struct {
|
||||||
rawWriter io.WriteCloser
|
rawWriter io.WriteCloser
|
||||||
|
ctx context.Context
|
||||||
|
|
||||||
rateBytes int
|
rateBytes int
|
||||||
|
|
||||||
@@ -18,9 +20,10 @@ type RateLimitWriter struct {
|
|||||||
before time.Time
|
before time.Time
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewRateLimitWriter(rawWriter io.WriteCloser, rateBytes int64) io.WriteCloser {
|
func NewRateLimitWriter(ctx context.Context, rawWriter io.WriteCloser, rateBytes int64) io.WriteCloser {
|
||||||
return &RateLimitWriter{
|
return &RateLimitWriter{
|
||||||
rawWriter: rawWriter,
|
rawWriter: rawWriter,
|
||||||
|
ctx: ctx,
|
||||||
rateBytes: types.Int(rateBytes),
|
rateBytes: types.Int(rateBytes),
|
||||||
before: time.Now(),
|
before: time.Now(),
|
||||||
}
|
}
|
||||||
@@ -71,6 +74,14 @@ func (this *RateLimitWriter) write(p []byte) (n int, err error) {
|
|||||||
n, err = this.rawWriter.Write(p)
|
n, err = this.rawWriter.Write(p)
|
||||||
|
|
||||||
if err == nil {
|
if err == nil {
|
||||||
|
select {
|
||||||
|
case <-this.ctx.Done():
|
||||||
|
err = io.EOF
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
this.written += n
|
this.written += n
|
||||||
|
|
||||||
if this.written >= this.rateBytes {
|
if this.written >= this.rateBytes {
|
||||||
|
|||||||
Reference in New Issue
Block a user