refractor: limiter builder

1. use map instead of sync.Map in defaultManger for concurrent read only
2. build limiterKey with struct
This commit is contained in:
zakuwaki 2023-07-24 11:39:51 +08:00
parent ab1412a0c9
commit 077b77a4a5

View File

@ -4,7 +4,6 @@ import (
"context" "context"
"fmt" "fmt"
"net" "net"
"sync"
"github.com/sagernet/sing-box/common/humanize" "github.com/sagernet/sing-box/common/humanize"
"github.com/sagernet/sing-box/log" "github.com/sagernet/sing-box/log"
@ -14,19 +13,24 @@ import (
) )
const ( const (
limiterTag = "tag" prefixTag = "tag"
limiterUser = "user" prefixUser = "user"
limiterInbound = "inbound" prefixInbound = "inbound"
) )
var _ Manager = (*defaultManager)(nil) var _ Manager = (*defaultManager)(nil)
type limiterKey struct {
Prefix string
Name string
}
type defaultManager struct { type defaultManager struct {
mp *sync.Map mp map[limiterKey]*limiter
} }
func WithDefault(ctx context.Context, logger log.ContextLogger, options []option.Limiter) context.Context { func WithDefault(ctx context.Context, logger log.ContextLogger, options []option.Limiter) context.Context {
m := &defaultManager{mp: &sync.Map{}} m := &defaultManager{mp: make(map[limiterKey]*limiter)}
for i, option := range options { for i, option := range options {
if err := m.createLimiter(ctx, option); err != nil { if err := m.createLimiter(ctx, option); err != nil {
logger.ErrorContext(ctx, fmt.Sprintf("id=%d, %s", i, err)) logger.ErrorContext(ctx, fmt.Sprintf("id=%d, %s", i, err))
@ -38,10 +42,6 @@ func WithDefault(ctx context.Context, logger log.ContextLogger, options []option
return service.ContextWith[Manager](ctx, m) return service.ContextWith[Manager](ctx, m)
} }
func buildKey(prefix string, tag string) string {
return fmt.Sprintf("%s|%s", prefix, tag)
}
func (m *defaultManager) createLimiter(ctx context.Context, option option.Limiter) (err error) { func (m *defaultManager) createLimiter(ctx context.Context, option option.Limiter) (err error) {
var download, upload uint64 var download, upload uint64
if len(option.Download) > 0 { if len(option.Download) > 0 {
@ -63,18 +63,18 @@ func (m *defaultManager) createLimiter(ctx context.Context, option option.Limite
valid := false valid := false
if len(option.Tag) > 0 { if len(option.Tag) > 0 {
valid = true valid = true
m.mp.Store(buildKey(limiterTag, option.Tag), l) m.mp[limiterKey{prefixTag, option.Tag}] = l
} }
if len(option.AuthUser) > 0 { if len(option.AuthUser) > 0 {
valid = true valid = true
for _, user := range option.AuthUser { for _, user := range option.AuthUser {
m.mp.Store(buildKey(limiterUser, user), l) m.mp[limiterKey{prefixUser, user}] = l
} }
} }
if len(option.Inbound) > 0 { if len(option.Inbound) > 0 {
valid = true valid = true
for _, inbound := range option.Inbound { for _, inbound := range option.Inbound {
m.mp.Store(buildKey(limiterInbound, inbound), l) m.mp[limiterKey{prefixInbound, inbound}] = l
} }
} }
if !valid { if !valid {
@ -84,16 +84,16 @@ func (m *defaultManager) createLimiter(ctx context.Context, option option.Limite
} }
func (m *defaultManager) LoadLimiters(tags []string, user, inbound string) (limiters []*limiter) { func (m *defaultManager) LoadLimiters(tags []string, user, inbound string) (limiters []*limiter) {
for _, t := range tags { for _, tag := range tags {
if v, ok := m.mp.Load(buildKey(limiterTag, t)); ok { if v, ok := m.mp[limiterKey{prefixTag, tag}]; ok {
limiters = append(limiters, v.(*limiter)) limiters = append(limiters, v)
} }
} }
if v, ok := m.mp.Load(buildKey(limiterUser, user)); ok { if v, ok := m.mp[limiterKey{prefixUser, user}]; ok {
limiters = append(limiters, v.(*limiter)) limiters = append(limiters, v)
} }
if v, ok := m.mp.Load(buildKey(limiterInbound, inbound)); ok { if v, ok := m.mp[limiterKey{prefixInbound, inbound}]; ok {
limiters = append(limiters, v.(*limiter)) limiters = append(limiters, v)
} }
return return
} }