feat: add quota limiter and quota API (#107)
This commit is contained in:
+59
-4
@@ -470,6 +470,7 @@ type ServiceConfig struct {
|
||||
Resolver string `yaml:",omitempty" json:"resolver,omitempty"`
|
||||
Hosts string `yaml:",omitempty" json:"hosts,omitempty"`
|
||||
Limiter string `yaml:",omitempty" json:"limiter,omitempty"`
|
||||
Quotas []string `yaml:",omitempty" json:"quotas,omitempty"`
|
||||
CLimiter string `yaml:"climiter,omitempty" json:"climiter,omitempty"`
|
||||
RLimiter string `yaml:"rlimiter,omitempty" json:"rlimiter,omitempty"`
|
||||
Logger string `yaml:",omitempty" json:"logger,omitempty"`
|
||||
@@ -485,10 +486,12 @@ type ServiceConfig struct {
|
||||
}
|
||||
|
||||
type ServiceStatus struct {
|
||||
CreateTime int64 `yaml:"createTime" json:"createTime"`
|
||||
State string `yaml:"state" json:"state"`
|
||||
Events []ServiceEvent `yaml:",omitempty" json:"events,omitempty"`
|
||||
Stats *ServiceStats `yaml:",omitempty" json:"stats,omitempty"`
|
||||
CreateTime int64 `yaml:"createTime" json:"createTime"`
|
||||
State string `yaml:"state" json:"state"`
|
||||
Events []ServiceEvent `yaml:",omitempty" json:"events,omitempty"`
|
||||
Stats *ServiceStats `yaml:",omitempty" json:"stats,omitempty"`
|
||||
StoppedByLimit bool `yaml:"stoppedByLimit,omitempty" json:"stoppedByLimit,omitempty"`
|
||||
Quotas []ServiceQuota `yaml:",omitempty" json:"quotas,omitempty"`
|
||||
}
|
||||
|
||||
type ServiceEvent struct {
|
||||
@@ -504,6 +507,57 @@ type ServiceStats struct {
|
||||
OutputBytes uint64 `yaml:"outputBytes" json:"outputBytes"`
|
||||
}
|
||||
|
||||
// QuotaConfig is a named cumulative traffic-volume limiter; the same name
|
||||
// referenced by several services (ServiceConfig.Quotas) shares one counter.
|
||||
type QuotaConfig struct {
|
||||
Name string `json:"name"`
|
||||
Limit string `yaml:",omitempty" json:"limit,omitempty"`
|
||||
StartsAt string `yaml:"startsAt,omitempty" json:"startsAt,omitempty"`
|
||||
ExpiresAt string `yaml:"expiresAt,omitempty" json:"expiresAt,omitempty"`
|
||||
Direction string `yaml:",omitempty" json:"direction,omitempty"`
|
||||
Flush string `yaml:",omitempty" json:"flush,omitempty"`
|
||||
Store *QuotaStoreConfig `yaml:",omitempty" json:"store,omitempty"`
|
||||
// read-only
|
||||
Status *QuotaStatus `yaml:",omitempty" json:"status,omitempty"`
|
||||
}
|
||||
|
||||
type QuotaStoreConfig struct {
|
||||
Type string `json:"type"` // "file" (default) | "redis" (stub, not implemented)
|
||||
File string `yaml:",omitempty" json:"file,omitempty"`
|
||||
Redis *QuotaRedisConfig `yaml:",omitempty" json:"redis,omitempty"`
|
||||
}
|
||||
|
||||
type QuotaRedisConfig struct {
|
||||
Addr string `json:"addr"`
|
||||
DB int `yaml:",omitempty" json:"db,omitempty"`
|
||||
Username string `yaml:",omitempty" json:"username,omitempty"`
|
||||
Password string `yaml:",omitempty" json:"password,omitempty"`
|
||||
Key string `yaml:",omitempty" json:"key,omitempty"`
|
||||
}
|
||||
|
||||
type QuotaStatus struct {
|
||||
Used uint64 `yaml:"used" json:"used"`
|
||||
Limit uint64 `yaml:"limit" json:"limit"`
|
||||
StartsAt int64 `yaml:"startsAt,omitempty" json:"startsAt,omitempty"`
|
||||
ExpiresAt int64 `yaml:"expiresAt,omitempty" json:"expiresAt,omitempty"`
|
||||
Active bool `yaml:"active" json:"active"`
|
||||
Expired bool `yaml:"expired,omitempty" json:"expired,omitempty"`
|
||||
Blocked bool `yaml:"blocked" json:"blocked"`
|
||||
Direction string `yaml:"direction,omitempty" json:"direction,omitempty"`
|
||||
}
|
||||
|
||||
type ServiceQuota struct {
|
||||
Name string `yaml:"name" json:"name"`
|
||||
Used uint64 `yaml:"used" json:"used"`
|
||||
Limit uint64 `yaml:"limit" json:"limit"`
|
||||
StartsAt int64 `yaml:"startsAt,omitempty" json:"startsAt,omitempty"`
|
||||
ExpiresAt int64 `yaml:"expiresAt,omitempty" json:"expiresAt,omitempty"`
|
||||
Active bool `yaml:"active" json:"active"`
|
||||
Expired bool `yaml:"expired,omitempty" json:"expired,omitempty"`
|
||||
Blocked bool `yaml:"blocked" json:"blocked"`
|
||||
Direction string `yaml:"direction,omitempty" json:"direction,omitempty"`
|
||||
}
|
||||
|
||||
type ChainConfig struct {
|
||||
Name string `json:"name"`
|
||||
Hops []*HopConfig `json:"hops"`
|
||||
@@ -573,6 +627,7 @@ type Config struct {
|
||||
SDs []*SDConfig `yaml:"sds,omitempty" json:"sds,omitempty"`
|
||||
Recorders []*RecorderConfig `yaml:",omitempty" json:"recorders,omitempty"`
|
||||
Limiters []*LimiterConfig `yaml:",omitempty" json:"limiters,omitempty"`
|
||||
Quotas []*QuotaConfig `yaml:",omitempty" json:"quotas,omitempty"`
|
||||
CLimiters []*LimiterConfig `yaml:"climiters,omitempty" json:"climiters,omitempty"`
|
||||
RLimiters []*LimiterConfig `yaml:"rlimiters,omitempty" json:"rlimiters,omitempty"`
|
||||
Observers []*ObserverConfig `yaml:",omitempty" json:"observers,omitempty"`
|
||||
|
||||
@@ -35,11 +35,13 @@ import (
|
||||
limiter_parser "github.com/go-gost/x/config/parsing/limiter"
|
||||
logger_parser "github.com/go-gost/x/config/parsing/logger"
|
||||
observer_parser "github.com/go-gost/x/config/parsing/observer"
|
||||
quota_parser "github.com/go-gost/x/config/parsing/quota"
|
||||
recorder_parser "github.com/go-gost/x/config/parsing/recorder"
|
||||
resolver_parser "github.com/go-gost/x/config/parsing/resolver"
|
||||
router_parser "github.com/go-gost/x/config/parsing/router"
|
||||
sd_parser "github.com/go-gost/x/config/parsing/sd"
|
||||
service_parser "github.com/go-gost/x/config/parsing/service"
|
||||
quota "github.com/go-gost/x/limiter/quota"
|
||||
"github.com/go-gost/x/registry"
|
||||
)
|
||||
|
||||
@@ -252,6 +254,16 @@ func register(cfg *config.Config) error {
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
var entries []named[*quota.Limiter]
|
||||
for _, c := range cfg.Quotas {
|
||||
entries = append(entries, named[*quota.Limiter]{c.Name, quota_parser.ParseQuotaLimiter(c)})
|
||||
}
|
||||
if err := registerGroup(entries, registry.QuotaLimiterRegistry()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
var entries []named[conn.ConnLimiter]
|
||||
for _, c := range cfg.CLimiters {
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/alecthomas/units"
|
||||
"github.com/go-gost/core/logger"
|
||||
"github.com/go-gost/x/config"
|
||||
xquota "github.com/go-gost/x/limiter/quota"
|
||||
)
|
||||
|
||||
// ParseQuotaLimiter is best-effort: malformed fields are logged and treated as unset.
|
||||
func ParseQuotaLimiter(cfg *config.QuotaConfig) *xquota.Limiter {
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
log := logger.Default().WithFields(map[string]any{
|
||||
"kind": "quota",
|
||||
"quota": cfg.Name,
|
||||
})
|
||||
|
||||
var limit uint64
|
||||
if s := strings.TrimSpace(cfg.Limit); s != "" && s != "0" {
|
||||
if v, err := units.ParseBase2Bytes(s); err != nil {
|
||||
log.Warnf("quota: parse limit %q: %v", s, err)
|
||||
} else if v > 0 {
|
||||
limit = uint64(v)
|
||||
}
|
||||
}
|
||||
|
||||
var dir xquota.Direction
|
||||
switch strings.ToLower(strings.TrimSpace(cfg.Direction)) {
|
||||
case "in":
|
||||
dir = xquota.DirectionIn
|
||||
case "out":
|
||||
dir = xquota.DirectionOut
|
||||
default:
|
||||
dir = xquota.DirectionTotal
|
||||
}
|
||||
|
||||
var flush time.Duration
|
||||
if s := strings.TrimSpace(cfg.Flush); s != "" {
|
||||
if d, err := time.ParseDuration(s); err != nil {
|
||||
log.Warnf("quota: parse flush %q: %v", s, err)
|
||||
} else {
|
||||
flush = d
|
||||
}
|
||||
}
|
||||
|
||||
return xquota.NewLimiter(cfg.Name, xquota.Options{
|
||||
Limit: limit,
|
||||
StartsAt: parseQuotaTime(cfg.StartsAt, log, "startsAt"),
|
||||
ExpiresAt: parseQuotaTime(cfg.ExpiresAt, log, "expiresAt"),
|
||||
Direction: dir,
|
||||
Flush: flush,
|
||||
Store: parseQuotaStore(cfg.Store, log),
|
||||
Logger: log,
|
||||
})
|
||||
}
|
||||
|
||||
func parseQuotaTime(s string, log logger.Logger, field string) time.Time {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return time.Time{}
|
||||
}
|
||||
t, err := time.Parse(time.RFC3339, s)
|
||||
if err != nil {
|
||||
log.Warnf("quota: parse %s %q: %v", field, s, err)
|
||||
return time.Time{}
|
||||
}
|
||||
return t
|
||||
}
|
||||
|
||||
func parseQuotaStore(sc *config.QuotaStoreConfig, log logger.Logger) xquota.Store {
|
||||
const defaultFile = "gost-quota.json"
|
||||
|
||||
if sc == nil || sc.Type == "" || strings.EqualFold(sc.Type, "file") {
|
||||
path := defaultFile
|
||||
if sc != nil && strings.TrimSpace(sc.File) != "" {
|
||||
path = sc.File
|
||||
}
|
||||
return xquota.NewFileStore(path)
|
||||
}
|
||||
|
||||
if strings.EqualFold(sc.Type, "redis") {
|
||||
log.Warnf("quota: redis store not implemented; counter will not be persisted")
|
||||
var rc config.QuotaRedisConfig
|
||||
if sc.Redis != nil {
|
||||
rc = *sc.Redis
|
||||
}
|
||||
return xquota.NewRedisStore(xquota.RedisConfig{
|
||||
Addr: rc.Addr,
|
||||
Username: rc.Username,
|
||||
Password: rc.Password,
|
||||
DB: rc.DB,
|
||||
Key: rc.Key,
|
||||
})
|
||||
}
|
||||
|
||||
log.Warnf("quota: unknown store type %q; using file store %q", sc.Type, defaultFile)
|
||||
return xquota.NewFileStore(defaultFile)
|
||||
}
|
||||
@@ -29,6 +29,7 @@ import (
|
||||
logger_parser "github.com/go-gost/x/config/parsing/logger"
|
||||
selector_parser "github.com/go-gost/x/config/parsing/selector"
|
||||
tls_util "github.com/go-gost/x/internal/util/tls"
|
||||
quota_wrapper "github.com/go-gost/x/limiter/quota/wrapper"
|
||||
cache_limiter "github.com/go-gost/x/limiter/traffic/cache"
|
||||
"github.com/go-gost/x/metadata"
|
||||
mdutil "github.com/go-gost/x/metadata/util"
|
||||
@@ -246,6 +247,10 @@ func ParseService(cfg *config.ServiceConfig) (service.Service, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for _, qname := range cfg.Quotas {
|
||||
ln = quota_wrapper.WrapListener(ln, strings.TrimSpace(qname))
|
||||
}
|
||||
|
||||
handlerLogger := serviceLogger.WithFields(map[string]any{
|
||||
"kind": "handler",
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user