ginuerzh
fa708f4b5f
Add package and exported symbol doc comments across conn, generator, limiter, and wrapper packages. Fix nil receiver guard in connLimitSingleGenerator.Limiter, add missing httpLoader.Close call, return context.Background instead of nil in serverConn.Context, return explicit error on connection limit exceeded, and simplify Allow logic.
71 lines
1.2 KiB
Go
71 lines
1.2 KiB
Go
package conn
|
|
|
|
import (
|
|
"sort"
|
|
"sync/atomic"
|
|
|
|
limiter "github.com/go-gost/core/limiter/conn"
|
|
)
|
|
|
|
type llimiter struct {
|
|
limit int
|
|
current int64
|
|
}
|
|
|
|
// NewLimiter creates a Limiter that allows up to n concurrent connections.
|
|
func NewLimiter(n int) limiter.Limiter {
|
|
return &llimiter{limit: n}
|
|
}
|
|
|
|
func (l *llimiter) Limit() int {
|
|
return l.limit
|
|
}
|
|
|
|
func (l *llimiter) Allow(n int) bool {
|
|
if n < 0 {
|
|
atomic.AddInt64(&l.current, int64(n))
|
|
return true
|
|
}
|
|
if atomic.AddInt64(&l.current, int64(n)) > int64(l.limit) {
|
|
atomic.AddInt64(&l.current, -int64(n))
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
type limiterGroup struct {
|
|
limiters []limiter.Limiter
|
|
}
|
|
|
|
func newLimiterGroup(limiters ...limiter.Limiter) *limiterGroup {
|
|
sort.Slice(limiters, func(i, j int) bool {
|
|
return limiters[i].Limit() < limiters[j].Limit()
|
|
})
|
|
return &limiterGroup{limiters: limiters}
|
|
}
|
|
|
|
func (l *limiterGroup) Allow(n int) (b bool) {
|
|
var i int
|
|
|
|
for i = range l.limiters {
|
|
if b = l.limiters[i].Allow(n); !b {
|
|
break
|
|
}
|
|
}
|
|
if !b && i > 0 && n > 0 {
|
|
for i := range l.limiters[:i] {
|
|
l.limiters[i].Allow(-n)
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func (l *limiterGroup) Limit() int {
|
|
if len(l.limiters) == 0 {
|
|
return 0
|
|
}
|
|
|
|
return l.limiters[0].Limit()
|
|
}
|