go-gost/core
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add Set for traffic Limiter

Browse Source
This commit is contained in:
ginuerzh
2022-12-19 19:34:13 +08:00
parent ceccf65ca1
commit ba87494cef
2 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
common/net/dialer/dialer.go
View File

@ -26,15 +26,16 @@ type NetDialer struct {
Timeout time.Duration Timeout time.Duration
DialFunc func(ctx context.Context, network, addr string) (net.Conn, error) DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)
Logger logger.Logger Logger logger.Logger
deadline time.Time
} }
func (d *NetDialer) Dial(ctx context.Context, network, addr string) (conn net.Conn, err error) { func (d *NetDialer) Dial(ctx context.Context, network, addr string) (conn net.Conn, err error) {
if d == nil { if d == nil {
d = DefaultNetDialer d = DefaultNetDialer
} }
if d.Timeout <= 0 {
d.Timeout = DefaultTimeout timeout := d.Timeout
if timeout <= 0 {
timeout = DefaultTimeout
} }
if d.DialFunc != nil { if d.DialFunc != nil {
@ -46,8 +47,8 @@ func (d *NetDialer) Dial(ctx context.Context, network, addr string) (conn net.Co
log = logger.Default() log = logger.Default()
} }
deadline := time.Now().Add(timeout)
ifces := strings.Split(d.Interface, ",") ifces := strings.Split(d.Interface, ",")
d.deadline = time.Now().Add(d.Timeout)
for _, ifce := range ifces { for _, ifce := range ifces {
strict := strings.HasSuffix(ifce, "!") strict := strings.HasSuffix(ifce, "!")
ifce = strings.TrimSuffix(ifce, "!") ifce = strings.TrimSuffix(ifce, "!")
@ -59,7 +60,7 @@ func (d *NetDialer) Dial(ctx context.Context, network, addr string) (conn net.Co
} }
for _, ifAddr := range ifAddrs { for _, ifAddr := range ifAddrs {
conn, err = d.dialOnce(ctx, network, addr, ifceName, ifAddr, log) conn, err = d.dialOnce(ctx, network, addr, ifceName, ifAddr, deadline, log)
if err == nil { if err == nil {
return return
} }
@ -72,7 +73,7 @@ func (d *NetDialer) Dial(ctx context.Context, network, addr string) (conn net.Co
return return
} }
if time.Until(d.deadline) < 0 { if time.Until(deadline) < 0 {
return return
} }
} }
@ -81,7 +82,7 @@ func (d *NetDialer) Dial(ctx context.Context, network, addr string) (conn net.Co
return return
} }
func (d *NetDialer) dialOnce(ctx context.Context, network, addr, ifceName string, ifAddr net.Addr, log logger.Logger) (net.Conn, error) { func (d *NetDialer) dialOnce(ctx context.Context, network, addr, ifceName string, ifAddr net.Addr, deadline time.Time, log logger.Logger) (net.Conn, error) {
if ifceName != "" { if ifceName != "" {
log.Debugf("interface: %s %v/%s", ifceName, ifAddr, network) log.Debugf("interface: %s %v/%s", ifceName, ifAddr, network)
} }
@ -125,7 +126,7 @@ func (d *NetDialer) dialOnce(ctx context.Context, network, addr, ifceName string
return nil, fmt.Errorf("dial: unsupported network %s", network) return nil, fmt.Errorf("dial: unsupported network %s", network)
} }
netd := net.Dialer{ netd := net.Dialer{
Deadline: d.deadline, Deadline: deadline,
LocalAddr: ifAddr, LocalAddr: ifAddr,
Control: func(network, address string, c syscall.RawConn) error { Control: func(network, address string, c syscall.RawConn) error {
return c.Control(func(fd uintptr) { return c.Control(func(fd uintptr) {

1
limiter/traffic/limiter.go
View File

@ -7,6 +7,7 @@ type Limiter interface {
// the returned value is less or equal to n. // the returned value is less or equal to n.
Wait(ctx context.Context, n int) int Wait(ctx context.Context, n int) int
Limit() int Limit() int
Set(n int)
} }
type TrafficLimiter interface { type TrafficLimiter interface {