feat: add quota limiter and quota API (#107)

This commit is contained in:
Usishchev Yury
2026-06-20 08:44:42 +02:00
committed by GitHub
parent b7d807b464
commit b23553d37a
16 changed files with 1343 additions and 4 deletions
+59 -4
View File
@@ -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"`
+12
View File
@@ -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 {
+104
View File
@@ -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)
}
+5
View File
@@ -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",
})