feat: add quota limiter and quota API (#107)
This commit is contained in:
Usishchev Yury
GitHub
@@ -145,6 +145,13 @@ func Register(r *gin.Engine, opts *Options) {
|
||||
config.PUT("/limiters/:limiter", updateLimiter)
|
||||
config.DELETE("/limiters/:limiter", deleteLimiter)
|
||||
|
||||
config.GET("/quotas", getQuotaList)
|
||||
config.GET("/quotas/:quota", getQuota)
|
||||
config.POST("/quotas", createQuota)
|
||||
config.PUT("/quotas/:quota", updateQuota)
|
||||
config.DELETE("/quotas/:quota", deleteQuota)
|
||||
config.POST("/quotas/:quota/reset", resetQuota)
|
||||
|
||||
config.GET("/climiters", getConnLimiterList)
|
||||
config.GET("/climiters/:limiter", getConnLimiter)
|
||||
config.POST("/climiters", createConnLimiter)
|
||||
|
||||
@@ -35,6 +35,29 @@ func fillServiceStatus(svc *config.ServiceConfig) {
|
||||
OutputBytes: st.Get(stats.KindOutputBytes),
|
||||
}
|
||||
}
|
||||
var stopped bool
|
||||
for _, qname := range svc.Quotas {
|
||||
lim := registry.QuotaLimiterRegistry().Get(qname)
|
||||
if lim == nil {
|
||||
continue
|
||||
}
|
||||
qs := lim.Snapshot()
|
||||
if qs.Blocked {
|
||||
stopped = true
|
||||
}
|
||||
svc.Status.Quotas = append(svc.Status.Quotas, config.ServiceQuota{
|
||||
Name: qname,
|
||||
Used: qs.Used,
|
||||
Limit: qs.Limit,
|
||||
StartsAt: qs.StartsAtUnix,
|
||||
ExpiresAt: qs.ExpiresAtUnix,
|
||||
Active: qs.Active,
|
||||
Expired: qs.Expired,
|
||||
Blocked: qs.Blocked,
|
||||
Direction: qs.Direction,
|
||||
})
|
||||
}
|
||||
svc.Status.StoppedByLimit = stopped
|
||||
for _, ev := range status.Events() {
|
||||
if !ev.Time.IsZero() {
|
||||
svc.Status.Events = append(svc.Status.Events, config.ServiceEvent{
|
||||
|
||||
@@ -0,0 +1,350 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-gost/x/config"
|
||||
parser "github.com/go-gost/x/config/parsing/quota"
|
||||
"github.com/go-gost/x/limiter/quota"
|
||||
"github.com/go-gost/x/registry"
|
||||
)
|
||||
|
||||
func fillQuotaStatus(q *config.QuotaConfig) {
|
||||
lim := registry.QuotaLimiterRegistry().Get(q.Name)
|
||||
if lim == nil {
|
||||
return
|
||||
}
|
||||
s := lim.Snapshot()
|
||||
q.Status = &config.QuotaStatus{
|
||||
Used: s.Used,
|
||||
Limit: s.Limit,
|
||||
StartsAt: s.StartsAtUnix,
|
||||
ExpiresAt: s.ExpiresAtUnix,
|
||||
Active: s.Active,
|
||||
Expired: s.Expired,
|
||||
Blocked: s.Blocked,
|
||||
Direction: s.Direction,
|
||||
}
|
||||
}
|
||||
|
||||
// swagger:parameters getQuotaListRequest
|
||||
type getQuotaListRequest struct {
|
||||
}
|
||||
|
||||
// successful operation.
|
||||
// swagger:response getQuotaListResponse
|
||||
type getQuotaListResponse struct {
|
||||
// in: body
|
||||
Data quotaList
|
||||
}
|
||||
|
||||
type quotaList struct {
|
||||
Count int `json:"count"`
|
||||
List []*config.QuotaConfig `json:"list"`
|
||||
}
|
||||
|
||||
func getQuotaList(ctx *gin.Context) {
|
||||
// swagger:route GET /config/quotas Quota getQuotaListRequest
|
||||
//
|
||||
// Get quota list.
|
||||
//
|
||||
// Security:
|
||||
// basicAuth: []
|
||||
//
|
||||
// Responses:
|
||||
// 200: getQuotaListResponse
|
||||
|
||||
var req getQuotaListRequest
|
||||
ctx.ShouldBindQuery(&req)
|
||||
|
||||
list := config.Global().Quotas
|
||||
for _, q := range list {
|
||||
if q == nil {
|
||||
continue
|
||||
}
|
||||
fillQuotaStatus(q)
|
||||
}
|
||||
|
||||
ctx.JSON(http.StatusOK, Response{
|
||||
Data: quotaList{
|
||||
Count: len(list),
|
||||
List: list,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// swagger:parameters getQuotaRequest
|
||||
type getQuotaRequest struct {
|
||||
// in: path
|
||||
// required: true
|
||||
Quota string `uri:"quota" json:"quota"`
|
||||
}
|
||||
|
||||
// successful operation.
|
||||
// swagger:response getQuotaResponse
|
||||
type getQuotaResponse struct {
|
||||
// in: body
|
||||
Data *config.QuotaConfig
|
||||
}
|
||||
|
||||
func getQuota(ctx *gin.Context) {
|
||||
// swagger:route GET /config/quotas/{quota} Quota getQuotaRequest
|
||||
//
|
||||
// Get quota.
|
||||
//
|
||||
// Security:
|
||||
// basicAuth: []
|
||||
//
|
||||
// Responses:
|
||||
// 200: getQuotaResponse
|
||||
|
||||
var req getQuotaRequest
|
||||
ctx.ShouldBindUri(&req)
|
||||
|
||||
var resp getQuotaResponse
|
||||
|
||||
for _, q := range config.Global().Quotas {
|
||||
if q == nil {
|
||||
continue
|
||||
}
|
||||
if q.Name == req.Quota {
|
||||
fillQuotaStatus(q)
|
||||
resp.Data = q
|
||||
}
|
||||
}
|
||||
|
||||
ctx.JSON(http.StatusOK, Response{
|
||||
Data: resp.Data,
|
||||
})
|
||||
}
|
||||
|
||||
// swagger:parameters createQuotaRequest
|
||||
type createQuotaRequest struct {
|
||||
// in: body
|
||||
Data config.QuotaConfig `json:"data"`
|
||||
}
|
||||
|
||||
// successful operation.
|
||||
// swagger:response createQuotaResponse
|
||||
type createQuotaResponse struct {
|
||||
Data Response
|
||||
}
|
||||
|
||||
func createQuota(ctx *gin.Context) {
|
||||
// swagger:route POST /config/quotas Quota createQuotaRequest
|
||||
//
|
||||
// Create a new quota, the name of the quota must be unique in quota list.
|
||||
//
|
||||
// Security:
|
||||
// basicAuth: []
|
||||
//
|
||||
// Responses:
|
||||
// 200: createQuotaResponse
|
||||
|
||||
var req createQuotaRequest
|
||||
ctx.ShouldBindJSON(&req.Data)
|
||||
|
||||
name := strings.TrimSpace(req.Data.Name)
|
||||
if name == "" {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeInvalid, "quota name is required"))
|
||||
return
|
||||
}
|
||||
req.Data.Name = name
|
||||
req.Data.Status = nil
|
||||
|
||||
if registry.QuotaLimiterRegistry().IsRegistered(name) {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeDup, fmt.Sprintf("quota %s already exists", name)))
|
||||
return
|
||||
}
|
||||
|
||||
v := parser.ParseQuotaLimiter(&req.Data)
|
||||
|
||||
if err := registry.QuotaLimiterRegistry().Register(name, v); err != nil {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeDup, fmt.Sprintf("quota %s already exists", name)))
|
||||
return
|
||||
}
|
||||
|
||||
config.OnUpdate(func(c *config.Config) error {
|
||||
c.Quotas = append(c.Quotas, &req.Data)
|
||||
return nil
|
||||
})
|
||||
|
||||
ctx.JSON(http.StatusOK, Response{
|
||||
Msg: "OK",
|
||||
})
|
||||
}
|
||||
|
||||
// swagger:parameters updateQuotaRequest
|
||||
type updateQuotaRequest struct {
|
||||
// in: path
|
||||
// required: true
|
||||
Quota string `uri:"quota" json:"quota"`
|
||||
// in: body
|
||||
Data config.QuotaConfig `json:"data"`
|
||||
}
|
||||
|
||||
// successful operation.
|
||||
// swagger:response updateQuotaResponse
|
||||
type updateQuotaResponse struct {
|
||||
Data Response
|
||||
}
|
||||
|
||||
func updateQuota(ctx *gin.Context) {
|
||||
// swagger:route PUT /config/quotas/{quota} Quota updateQuotaRequest
|
||||
//
|
||||
// Update quota by name, the quota must already exist. The cumulative counter
|
||||
// is preserved across the update as long as the window is unchanged.
|
||||
//
|
||||
// Security:
|
||||
// basicAuth: []
|
||||
//
|
||||
// Responses:
|
||||
// 200: updateQuotaResponse
|
||||
|
||||
var req updateQuotaRequest
|
||||
ctx.ShouldBindUri(&req)
|
||||
ctx.ShouldBindJSON(&req.Data)
|
||||
|
||||
name := strings.TrimSpace(req.Quota)
|
||||
|
||||
if !registry.QuotaLimiterRegistry().IsRegistered(name) {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeNotFound, fmt.Sprintf("quota %s not found", name)))
|
||||
return
|
||||
}
|
||||
|
||||
req.Data.Name = name
|
||||
req.Data.Status = nil
|
||||
|
||||
v := parser.ParseQuotaLimiter(&req.Data)
|
||||
|
||||
registry.QuotaLimiterRegistry().Unregister(name)
|
||||
|
||||
if err := registry.QuotaLimiterRegistry().Register(name, v); err != nil {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeDup, fmt.Sprintf("quota %s already exists", name)))
|
||||
return
|
||||
}
|
||||
|
||||
config.OnUpdate(func(c *config.Config) error {
|
||||
for i := range c.Quotas {
|
||||
if c.Quotas[i].Name == name {
|
||||
c.Quotas[i] = &req.Data
|
||||
break
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
ctx.JSON(http.StatusOK, Response{
|
||||
Msg: "OK",
|
||||
})
|
||||
}
|
||||
|
||||
// swagger:parameters deleteQuotaRequest
|
||||
type deleteQuotaRequest struct {
|
||||
// in: path
|
||||
// required: true
|
||||
Quota string `uri:"quota" json:"quota"`
|
||||
}
|
||||
|
||||
// successful operation.
|
||||
// swagger:response deleteQuotaResponse
|
||||
type deleteQuotaResponse struct {
|
||||
Data Response
|
||||
}
|
||||
|
||||
func deleteQuota(ctx *gin.Context) {
|
||||
// swagger:route DELETE /config/quotas/{quota} Quota deleteQuotaRequest
|
||||
//
|
||||
// Delete quota by name. Services referencing it stop being limited (fail-open).
|
||||
//
|
||||
// Security:
|
||||
// basicAuth: []
|
||||
//
|
||||
// Responses:
|
||||
// 200: deleteQuotaResponse
|
||||
|
||||
var req deleteQuotaRequest
|
||||
ctx.ShouldBindUri(&req)
|
||||
|
||||
name := strings.TrimSpace(req.Quota)
|
||||
|
||||
if !registry.QuotaLimiterRegistry().IsRegistered(name) {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeNotFound, fmt.Sprintf("quota %s not found", name)))
|
||||
return
|
||||
}
|
||||
registry.QuotaLimiterRegistry().Unregister(name)
|
||||
|
||||
config.OnUpdate(func(c *config.Config) error {
|
||||
quotas := c.Quotas
|
||||
c.Quotas = nil
|
||||
for _, q := range quotas {
|
||||
if q.Name == name {
|
||||
continue
|
||||
}
|
||||
c.Quotas = append(c.Quotas, q)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
ctx.JSON(http.StatusOK, Response{
|
||||
Msg: "OK",
|
||||
})
|
||||
}
|
||||
|
||||
// swagger:parameters resetQuotaRequest
|
||||
type resetQuotaRequest struct {
|
||||
// in: path
|
||||
// required: true
|
||||
Quota string `uri:"quota" json:"quota"`
|
||||
// in: body
|
||||
Data resetQuotaData `json:"data"`
|
||||
}
|
||||
|
||||
type resetQuotaData struct {
|
||||
Used *uint64 `json:"used,omitempty"`
|
||||
}
|
||||
|
||||
// successful operation.
|
||||
// swagger:response resetQuotaResponse
|
||||
type resetQuotaResponse struct {
|
||||
Data Response
|
||||
}
|
||||
|
||||
func resetQuota(ctx *gin.Context) {
|
||||
// swagger:route POST /config/quotas/{quota}/reset Quota resetQuotaRequest
|
||||
//
|
||||
// Overwrite the cumulative counter of a quota (defaults to 0). The new value
|
||||
// is applied immediately and persisted; if the quota was blocking, services
|
||||
// resume once the value is below the limit.
|
||||
//
|
||||
// Security:
|
||||
// basicAuth: []
|
||||
//
|
||||
// Responses:
|
||||
// 200: resetQuotaResponse
|
||||
|
||||
var req resetQuotaRequest
|
||||
ctx.ShouldBindUri(&req)
|
||||
ctx.ShouldBindJSON(&req.Data)
|
||||
|
||||
name := strings.TrimSpace(req.Quota)
|
||||
|
||||
lim := registry.QuotaLimiterRegistry().Get(name)
|
||||
if lim == nil {
|
||||
writeError(ctx, NewError(http.StatusBadRequest, ErrCodeNotFound, fmt.Sprintf("quota %s not found", name)))
|
||||
return
|
||||
}
|
||||
|
||||
used := uint64(0)
|
||||
if req.Data.Used != nil {
|
||||
used = *req.Data.Used
|
||||
}
|
||||
lim.Update(quota.Update{Used: &used})
|
||||
|
||||
ctx.JSON(http.StatusOK, Response{
|
||||
Msg: "OK",
|
||||
})
|
||||
}
|
||||
+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",
|
||||
})
|
||||
|
||||
@@ -0,0 +1,326 @@
|
||||
// Package quota implements a persisted, long-term traffic-volume limiter that
|
||||
// can be shared across services by name. Unlike the rate limiters in
|
||||
// limiter/traffic and limiter/rate, it accumulates total bytes within a window
|
||||
// [startsAt, expiresAt) and stops the referencing service(s) once a byte limit
|
||||
// is reached. Enforcement is fail-open outside the window.
|
||||
package quota
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/logger"
|
||||
xlogger "github.com/go-gost/x/logger"
|
||||
)
|
||||
|
||||
var ErrQuotaExceeded = errors.New("quota: traffic limit reached")
|
||||
|
||||
const defaultFlushInterval = 10 * time.Second
|
||||
|
||||
type Direction int
|
||||
|
||||
const (
|
||||
DirectionTotal Direction = iota
|
||||
DirectionIn
|
||||
DirectionOut
|
||||
)
|
||||
|
||||
func (d Direction) String() string {
|
||||
switch d {
|
||||
case DirectionIn:
|
||||
return "in"
|
||||
case DirectionOut:
|
||||
return "out"
|
||||
default:
|
||||
return "total"
|
||||
}
|
||||
}
|
||||
|
||||
// Options seeds a Limiter. Limit and the window are config-authoritative; a
|
||||
// persisted counter is restored only for a matching window (see NewLimiter).
|
||||
type Options struct {
|
||||
Limit uint64
|
||||
StartsAt time.Time
|
||||
ExpiresAt time.Time
|
||||
Direction Direction
|
||||
Flush time.Duration
|
||||
Store Store
|
||||
Logger logger.Logger
|
||||
}
|
||||
|
||||
// Update overwrites runtime state; a nil field is left unchanged.
|
||||
type Update struct {
|
||||
Used *uint64
|
||||
Limit *uint64
|
||||
StartsAt *time.Time
|
||||
ExpiresAt *time.Time
|
||||
}
|
||||
|
||||
type Snapshot struct {
|
||||
Used uint64
|
||||
Limit uint64
|
||||
StartsAtUnix int64
|
||||
ExpiresAtUnix int64
|
||||
Active bool
|
||||
Expired bool
|
||||
Blocked bool
|
||||
Direction string
|
||||
}
|
||||
|
||||
type Limiter struct {
|
||||
name string
|
||||
direction Direction
|
||||
flush time.Duration
|
||||
store Store
|
||||
log logger.Logger
|
||||
|
||||
used atomic.Uint64
|
||||
limit atomic.Uint64
|
||||
startsAt atomic.Int64 // unixnano; 0 = unset
|
||||
expiresAt atomic.Int64 // unixnano; 0 = unset
|
||||
blocked atomic.Bool
|
||||
dirty atomic.Bool
|
||||
closed atomic.Bool
|
||||
|
||||
mu sync.Mutex
|
||||
waitCh chan struct{} // closed+replaced to broadcast a state change
|
||||
|
||||
closeOnce sync.Once
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func NewLimiter(name string, opts Options) *Limiter {
|
||||
l := &Limiter{
|
||||
name: name,
|
||||
direction: opts.Direction,
|
||||
flush: opts.Flush,
|
||||
store: opts.Store,
|
||||
log: opts.Logger,
|
||||
waitCh: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
if l.flush <= 0 {
|
||||
l.flush = defaultFlushInterval
|
||||
}
|
||||
if l.log == nil {
|
||||
l.log = xlogger.Nop()
|
||||
}
|
||||
|
||||
l.limit.Store(opts.Limit)
|
||||
sa := unixNanoOrZero(opts.StartsAt)
|
||||
ea := unixNanoOrZero(opts.ExpiresAt)
|
||||
l.startsAt.Store(sa)
|
||||
l.expiresAt.Store(ea)
|
||||
|
||||
// Restore the counter only within the same window: a changed window (a new
|
||||
// period pushed via config) starts fresh.
|
||||
if l.store != nil {
|
||||
if rec, ok, err := l.store.Load(name); err != nil {
|
||||
l.log.Warnf("quota: load %s: %v", name, err)
|
||||
} else if ok && rec.StartsAt == sa && rec.ExpiresAt == ea {
|
||||
l.used.Store(rec.Used)
|
||||
}
|
||||
}
|
||||
|
||||
l.blocked.Store(l.enforcing(time.Now()) && l.used.Load() >= l.limit.Load())
|
||||
go l.run()
|
||||
return l
|
||||
}
|
||||
|
||||
func (l *Limiter) active(now time.Time) bool {
|
||||
n := now.UnixNano()
|
||||
if sa := l.startsAt.Load(); sa != 0 && n < sa {
|
||||
return false
|
||||
}
|
||||
if ea := l.expiresAt.Load(); ea != 0 && n >= ea {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (l *Limiter) enforcing(now time.Time) bool {
|
||||
return l.limit.Load() > 0 && l.active(now)
|
||||
}
|
||||
|
||||
func (l *Limiter) Blocked() bool {
|
||||
if l.closed.Load() || !l.blocked.Load() {
|
||||
return false
|
||||
}
|
||||
return l.enforcing(time.Now())
|
||||
}
|
||||
|
||||
func (l *Limiter) AddIn(n int) { l.add(int64(n), DirectionIn) }
|
||||
func (l *Limiter) AddOut(n int) { l.add(int64(n), DirectionOut) }
|
||||
|
||||
func (l *Limiter) add(n int64, dir Direction) {
|
||||
if n <= 0 || l.closed.Load() {
|
||||
return
|
||||
}
|
||||
switch l.direction {
|
||||
case DirectionIn:
|
||||
if dir != DirectionIn {
|
||||
return
|
||||
}
|
||||
case DirectionOut:
|
||||
if dir != DirectionOut {
|
||||
return
|
||||
}
|
||||
}
|
||||
if !l.active(time.Now()) {
|
||||
return
|
||||
}
|
||||
used := l.used.Add(uint64(n))
|
||||
l.dirty.Store(true)
|
||||
if lim := l.limit.Load(); lim > 0 && used >= lim {
|
||||
l.block()
|
||||
}
|
||||
}
|
||||
|
||||
// block is hot-path: no lock, no I/O, no wakeup (parking only happens on the
|
||||
// next Accept).
|
||||
func (l *Limiter) block() {
|
||||
if l.blocked.CompareAndSwap(false, true) {
|
||||
l.log.Warnf("quota: %s reached limit %d bytes, stopping", l.name, l.limit.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// WaitChan must be fetched before checking Blocked to avoid a missed wakeup.
|
||||
func (l *Limiter) WaitChan() <-chan struct{} {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.waitCh
|
||||
}
|
||||
|
||||
func (l *Limiter) notify() {
|
||||
l.mu.Lock()
|
||||
close(l.waitCh)
|
||||
l.waitCh = make(chan struct{})
|
||||
l.mu.Unlock()
|
||||
}
|
||||
|
||||
func (l *Limiter) Update(u Update) {
|
||||
if u.Used != nil {
|
||||
l.used.Store(*u.Used)
|
||||
}
|
||||
if u.Limit != nil {
|
||||
l.limit.Store(*u.Limit)
|
||||
}
|
||||
if u.StartsAt != nil {
|
||||
l.startsAt.Store(unixNanoOrZero(*u.StartsAt))
|
||||
}
|
||||
if u.ExpiresAt != nil {
|
||||
l.expiresAt.Store(unixNanoOrZero(*u.ExpiresAt))
|
||||
}
|
||||
l.reevaluate()
|
||||
}
|
||||
|
||||
func (l *Limiter) reevaluate() {
|
||||
now := time.Now()
|
||||
shouldBlock := l.enforcing(now) && l.used.Load() >= l.limit.Load()
|
||||
l.blocked.Store(shouldBlock)
|
||||
if !shouldBlock {
|
||||
l.notify()
|
||||
}
|
||||
l.flushNow()
|
||||
}
|
||||
|
||||
func (l *Limiter) Snapshot() Snapshot {
|
||||
now := time.Now()
|
||||
ea := l.expiresAt.Load()
|
||||
return Snapshot{
|
||||
Used: l.used.Load(),
|
||||
Limit: l.limit.Load(),
|
||||
StartsAtUnix: nanoToSec(l.startsAt.Load()),
|
||||
ExpiresAtUnix: nanoToSec(ea),
|
||||
Active: l.active(now),
|
||||
Expired: ea != 0 && now.UnixNano() >= ea,
|
||||
Blocked: l.Blocked(),
|
||||
Direction: l.direction.String(),
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Limiter) run() {
|
||||
flushT := time.NewTicker(l.flush)
|
||||
defer flushT.Stop()
|
||||
boundT := time.NewTicker(time.Second)
|
||||
defer boundT.Stop()
|
||||
|
||||
wasActive := l.active(time.Now())
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-l.done:
|
||||
return
|
||||
|
||||
case <-flushT.C:
|
||||
if l.dirty.Swap(false) {
|
||||
l.persist()
|
||||
}
|
||||
|
||||
case <-boundT.C:
|
||||
now := time.Now()
|
||||
act := l.active(now)
|
||||
if act == wasActive {
|
||||
continue
|
||||
}
|
||||
wasActive = act
|
||||
if act {
|
||||
l.blocked.Store(l.limit.Load() > 0 && l.used.Load() >= l.limit.Load())
|
||||
} else {
|
||||
l.blocked.Store(false)
|
||||
}
|
||||
l.notify()
|
||||
l.persist()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (l *Limiter) flushNow() {
|
||||
l.dirty.Store(false)
|
||||
l.persist()
|
||||
}
|
||||
|
||||
func (l *Limiter) persist() {
|
||||
if l.store == nil {
|
||||
return
|
||||
}
|
||||
rec := Record{
|
||||
Used: l.used.Load(),
|
||||
Limit: l.limit.Load(),
|
||||
StartsAt: l.startsAt.Load(),
|
||||
ExpiresAt: l.expiresAt.Load(),
|
||||
UpdatedAt: time.Now().UnixNano(),
|
||||
}
|
||||
if err := l.store.Save(l.name, rec); err != nil {
|
||||
l.log.Warnf("quota: save %s: %v", l.name, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Close makes the limiter inert (Blocked false, counting a no-op) so deleting a
|
||||
// shared quota releases the referencing services instead of tearing them down.
|
||||
// Called by the registry on Unregister.
|
||||
func (l *Limiter) Close() error {
|
||||
l.closeOnce.Do(func() {
|
||||
l.closed.Store(true)
|
||||
close(l.done)
|
||||
l.flushNow()
|
||||
l.notify()
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func unixNanoOrZero(t time.Time) int64 {
|
||||
if t.IsZero() {
|
||||
return 0
|
||||
}
|
||||
return t.UnixNano()
|
||||
}
|
||||
|
||||
func nanoToSec(n int64) int64 {
|
||||
if n == 0 {
|
||||
return 0
|
||||
}
|
||||
return n / int64(time.Second)
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newTestLimiter(t *testing.T, opts Options) *Limiter {
|
||||
t.Helper()
|
||||
l := NewLimiter("test", opts)
|
||||
t.Cleanup(func() { l.Close() })
|
||||
return l
|
||||
}
|
||||
|
||||
func TestLimiterBlocksAtLimit(t *testing.T) {
|
||||
l := newTestLimiter(t, Options{Limit: 100})
|
||||
|
||||
l.AddIn(60)
|
||||
if l.Blocked() {
|
||||
t.Fatalf("should not block at 60/100")
|
||||
}
|
||||
l.AddOut(40)
|
||||
if !l.Blocked() {
|
||||
t.Fatalf("should block at 100/100")
|
||||
}
|
||||
if s := l.Snapshot(); !s.Blocked || s.Used != 100 {
|
||||
t.Fatalf("snapshot=%+v, want blocked used=100", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterUnblockOnOverwrite(t *testing.T) {
|
||||
l := newTestLimiter(t, Options{Limit: 100})
|
||||
l.AddIn(150)
|
||||
if !l.Blocked() {
|
||||
t.Fatalf("expected blocked")
|
||||
}
|
||||
|
||||
zero := uint64(0)
|
||||
l.Update(Update{Used: &zero})
|
||||
|
||||
if l.Blocked() {
|
||||
t.Fatalf("expected unblocked after reset")
|
||||
}
|
||||
if got := l.Snapshot().Used; got != 0 {
|
||||
t.Fatalf("used=%d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterDirection(t *testing.T) {
|
||||
l := newTestLimiter(t, Options{Limit: 100, Direction: DirectionIn})
|
||||
|
||||
l.AddOut(500)
|
||||
if u := l.Snapshot().Used; u != 0 {
|
||||
t.Fatalf("outbound counted under DirectionIn: used=%d", u)
|
||||
}
|
||||
l.AddIn(100)
|
||||
if !l.Blocked() {
|
||||
t.Fatalf("expected blocked after inbound reaches limit")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterWindowNotStarted(t *testing.T) {
|
||||
l := newTestLimiter(t, Options{Limit: 100, StartsAt: time.Now().Add(time.Hour)})
|
||||
|
||||
l.AddIn(500)
|
||||
s := l.Snapshot()
|
||||
if s.Used != 0 {
|
||||
t.Fatalf("counted before startsAt: used=%d", s.Used)
|
||||
}
|
||||
if s.Active {
|
||||
t.Fatalf("should be inactive before startsAt")
|
||||
}
|
||||
if l.Blocked() {
|
||||
t.Fatalf("should not block before startsAt")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterFailOpenAfterExpiry(t *testing.T) {
|
||||
l := newTestLimiter(t, Options{Limit: 100, ExpiresAt: time.Now().Add(-time.Hour)})
|
||||
|
||||
l.AddIn(500)
|
||||
s := l.Snapshot()
|
||||
if s.Used != 0 {
|
||||
t.Fatalf("counted after expiry: used=%d", s.Used)
|
||||
}
|
||||
if !s.Expired {
|
||||
t.Fatalf("snapshot should report expired")
|
||||
}
|
||||
if l.Blocked() {
|
||||
t.Fatalf("should be fail-open (unblocked) after expiry")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterPersistsToDisk(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "quota.json")
|
||||
store := NewFileStore(path)
|
||||
|
||||
l := NewLimiter("svc", Options{Limit: 1000, Store: store})
|
||||
l.AddIn(300)
|
||||
l.AddOut(200)
|
||||
l.Close()
|
||||
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read store file: %v", err)
|
||||
}
|
||||
var data map[string]Record
|
||||
if err := json.Unmarshal(b, &data); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if rec := data["svc"]; rec.Used != 500 {
|
||||
t.Fatalf("persisted used=%d, want 500", rec.Used)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterRestoresFromDisk(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "quota.json")
|
||||
if err := os.WriteFile(path, []byte(`{"svc":{"used":777,"limit":1000}}`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
store := NewFileStore(path)
|
||||
l := NewLimiter("svc", Options{Limit: 1000, Store: store})
|
||||
t.Cleanup(func() { l.Close() })
|
||||
|
||||
if got := l.Snapshot().Used; got != 777 {
|
||||
t.Fatalf("restored used=%d, want 777", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitChanNotifiesOnReset(t *testing.T) {
|
||||
l := newTestLimiter(t, Options{Limit: 100})
|
||||
l.AddIn(200)
|
||||
if !l.Blocked() {
|
||||
t.Fatal("expected blocked")
|
||||
}
|
||||
|
||||
ch := l.WaitChan()
|
||||
zero := uint64(0)
|
||||
l.Update(Update{Used: &zero})
|
||||
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("WaitChan not notified after reset")
|
||||
}
|
||||
if l.Blocked() {
|
||||
t.Fatal("expected unblocked after reset")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClosedLimiterIsInert(t *testing.T) {
|
||||
l := NewLimiter("c", Options{Limit: 100})
|
||||
l.AddIn(200)
|
||||
if !l.Blocked() {
|
||||
t.Fatal("expected blocked")
|
||||
}
|
||||
|
||||
ch := l.WaitChan()
|
||||
l.Close()
|
||||
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("WaitChan not notified on close")
|
||||
}
|
||||
if l.Blocked() {
|
||||
t.Fatal("closed limiter must be inert (not blocking)")
|
||||
}
|
||||
l.AddIn(1000)
|
||||
if u := l.Snapshot().Used; u != 200 {
|
||||
t.Fatalf("closed limiter counted: used=%d, want 200", u)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterResetsOnWindowChange(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "quota.json")
|
||||
store := NewFileStore(path)
|
||||
|
||||
t1 := time.Now().Add(time.Hour)
|
||||
l1 := NewLimiter("svc", Options{Limit: 1000, ExpiresAt: t1, Store: store})
|
||||
l1.AddIn(500)
|
||||
l1.Close()
|
||||
|
||||
t2 := time.Now().Add(2 * time.Hour)
|
||||
l2 := NewLimiter("svc", Options{Limit: 1000, ExpiresAt: t2, Store: store})
|
||||
t.Cleanup(func() { l2.Close() })
|
||||
|
||||
if u := l2.Snapshot().Used; u != 0 {
|
||||
t.Fatalf("counter not reset on window change: used=%d, want 0", u)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package quota
|
||||
|
||||
type Record struct {
|
||||
Used uint64 `json:"used"`
|
||||
Limit uint64 `json:"limit"`
|
||||
StartsAt int64 `json:"startsAt"` // unixnano; 0 = unset
|
||||
ExpiresAt int64 `json:"expiresAt"` // unixnano; 0 = unset
|
||||
UpdatedAt int64 `json:"updatedAt"`
|
||||
}
|
||||
|
||||
// Store persists per-name quota records. Implementations must be concurrency-safe.
|
||||
type Store interface {
|
||||
// Load returns ok=false when no record exists yet.
|
||||
Load(name string) (rec Record, ok bool, err error)
|
||||
Save(name string, rec Record) error
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// fileStore persists all records into one JSON file. A single instance is shared
|
||||
// per path so concurrent per-name saves merge instead of clobbering each other.
|
||||
type fileStore struct {
|
||||
path string
|
||||
mu sync.Mutex
|
||||
data map[string]Record
|
||||
}
|
||||
|
||||
var fileStores sync.Map // path -> *fileStore
|
||||
|
||||
func NewFileStore(path string) Store {
|
||||
if v, ok := fileStores.Load(path); ok {
|
||||
return v.(*fileStore)
|
||||
}
|
||||
fs := &fileStore{path: path, data: make(map[string]Record)}
|
||||
fs.load()
|
||||
actual, _ := fileStores.LoadOrStore(path, fs)
|
||||
return actual.(*fileStore)
|
||||
}
|
||||
|
||||
func (fs *fileStore) load() {
|
||||
b, err := os.ReadFile(fs.path)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
data := make(map[string]Record)
|
||||
if json.Unmarshal(b, &data) == nil {
|
||||
fs.data = data
|
||||
}
|
||||
}
|
||||
|
||||
func (fs *fileStore) Load(name string) (Record, bool, error) {
|
||||
fs.mu.Lock()
|
||||
defer fs.mu.Unlock()
|
||||
rec, ok := fs.data[name]
|
||||
return rec, ok, nil
|
||||
}
|
||||
|
||||
func (fs *fileStore) Save(name string, rec Record) error {
|
||||
fs.mu.Lock()
|
||||
defer fs.mu.Unlock()
|
||||
fs.data[name] = rec
|
||||
return fs.writeLocked()
|
||||
}
|
||||
|
||||
func (fs *fileStore) writeLocked() error {
|
||||
b, err := json.MarshalIndent(fs.data, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tmp := fs.path + ".tmp"
|
||||
if err := os.WriteFile(tmp, b, 0o644); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, fs.path)
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package quota
|
||||
|
||||
type RedisConfig struct {
|
||||
Addr string
|
||||
Username string
|
||||
Password string
|
||||
DB int
|
||||
Key string
|
||||
}
|
||||
|
||||
// redisStore is a no-op placeholder: a redis-configured quota behaves as
|
||||
// in-memory only until the backend is implemented.
|
||||
//
|
||||
// TODO: implement using github.com/go-redis/redis/v8 (already a dependency).
|
||||
// Suggested schema: HSET <Key> <name> <json(Record)> for Save, HGET for Load.
|
||||
// See internal/loader/redis.go and recorder/redis.go for the patterns.
|
||||
type redisStore struct {
|
||||
cfg RedisConfig
|
||||
}
|
||||
|
||||
func NewRedisStore(cfg RedisConfig) Store {
|
||||
return &redisStore{cfg: cfg}
|
||||
}
|
||||
|
||||
func (s *redisStore) Load(name string) (Record, bool, error) {
|
||||
return Record{}, false, nil
|
||||
}
|
||||
|
||||
func (s *redisStore) Save(name string, rec Record) error {
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package wrapper
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"syscall"
|
||||
|
||||
"github.com/go-gost/x/ctx"
|
||||
xio "github.com/go-gost/x/internal/io"
|
||||
"github.com/go-gost/x/limiter/quota"
|
||||
"github.com/go-gost/x/registry"
|
||||
)
|
||||
|
||||
type quotaConn struct {
|
||||
net.Conn
|
||||
name string
|
||||
}
|
||||
|
||||
func WrapConn(c net.Conn, name string) net.Conn {
|
||||
if c == nil || name == "" {
|
||||
return c
|
||||
}
|
||||
return "aConn{Conn: c, name: name}
|
||||
}
|
||||
|
||||
func (c *quotaConn) Read(b []byte) (n int, err error) {
|
||||
lim := registry.QuotaLimiterRegistry().Get(c.name)
|
||||
if lim != nil && lim.Blocked() {
|
||||
return 0, quota.ErrQuotaExceeded
|
||||
}
|
||||
n, err = c.Conn.Read(b)
|
||||
if lim != nil && n > 0 {
|
||||
lim.AddIn(n)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (c *quotaConn) Write(b []byte) (n int, err error) {
|
||||
lim := registry.QuotaLimiterRegistry().Get(c.name)
|
||||
if lim != nil && lim.Blocked() {
|
||||
return 0, quota.ErrQuotaExceeded
|
||||
}
|
||||
n, err = c.Conn.Write(b)
|
||||
if lim != nil && n > 0 {
|
||||
lim.AddOut(n)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Forward optional capabilities so wrapping does not hide them from handlers.
|
||||
|
||||
func (c *quotaConn) SyscallConn() (syscall.RawConn, error) {
|
||||
if sc, ok := c.Conn.(syscall.Conn); ok {
|
||||
return sc.SyscallConn()
|
||||
}
|
||||
return nil, xio.ErrUnsupported
|
||||
}
|
||||
|
||||
func (c *quotaConn) CloseRead() error {
|
||||
if sc, ok := c.Conn.(xio.CloseRead); ok {
|
||||
return sc.CloseRead()
|
||||
}
|
||||
return xio.ErrUnsupported
|
||||
}
|
||||
|
||||
func (c *quotaConn) CloseWrite() error {
|
||||
if sc, ok := c.Conn.(xio.CloseWrite); ok {
|
||||
return sc.CloseWrite()
|
||||
}
|
||||
return xio.ErrUnsupported
|
||||
}
|
||||
|
||||
func (c *quotaConn) Context() context.Context {
|
||||
if cc, ok := c.Conn.(ctx.Context); ok {
|
||||
return cc.Context()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
// Package wrapper adapts a named quota onto a listener and its connections,
|
||||
// resolving the quota.Limiter by name from the registry on every call so
|
||||
// create/update/delete takes effect live and several services can share one.
|
||||
package wrapper
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/go-gost/core/listener"
|
||||
"github.com/go-gost/x/registry"
|
||||
)
|
||||
|
||||
type quotaListener struct {
|
||||
listener.Listener
|
||||
name string
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func WrapListener(ln listener.Listener, name string) listener.Listener {
|
||||
if ln == nil || name == "" {
|
||||
return ln
|
||||
}
|
||||
return "aListener{Listener: ln, name: name, done: make(chan struct{})}
|
||||
}
|
||||
|
||||
func (ln *quotaListener) Accept() (net.Conn, error) {
|
||||
for {
|
||||
lim := registry.QuotaLimiterRegistry().Get(ln.name)
|
||||
if lim == nil {
|
||||
break
|
||||
}
|
||||
ch := lim.WaitChan() // before Blocked, to avoid a missed wakeup
|
||||
if !lim.Blocked() {
|
||||
break
|
||||
}
|
||||
// Park (not error) so the service's accept loop survives; ln.done wakes
|
||||
// it on listener close.
|
||||
select {
|
||||
case <-ch:
|
||||
case <-ln.done:
|
||||
return nil, net.ErrClosed
|
||||
}
|
||||
}
|
||||
|
||||
c, err := ln.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return WrapConn(c, ln.name), nil
|
||||
}
|
||||
|
||||
func (ln *quotaListener) Close() error {
|
||||
// Wake a parked Accept, but do not close the (possibly shared) limiter.
|
||||
ln.once.Do(func() { close(ln.done) })
|
||||
return ln.Listener.Close()
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package registry
|
||||
|
||||
import "github.com/go-gost/x/limiter/quota"
|
||||
|
||||
// quotaLimiterRegistry holds shared quota limiters; Unregister closes them.
|
||||
type quotaLimiterRegistry struct {
|
||||
registry[*quota.Limiter]
|
||||
}
|
||||
@@ -23,6 +23,7 @@ import (
|
||||
"github.com/go-gost/core/router"
|
||||
"github.com/go-gost/core/sd"
|
||||
"github.com/go-gost/core/service"
|
||||
"github.com/go-gost/x/limiter/quota"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -49,6 +50,7 @@ var (
|
||||
trafficLimiterReg reg.Registry[traffic.TrafficLimiter] = new(trafficLimiterRegistry)
|
||||
connLimiterReg reg.Registry[conn.ConnLimiter] = new(connLimiterRegistry)
|
||||
rateLimiterReg reg.Registry[rate.RateLimiter] = new(rateLimiterRegistry)
|
||||
quotaLimiterReg reg.Registry[*quota.Limiter] = new(quotaLimiterRegistry)
|
||||
|
||||
ingressReg reg.Registry[ingress.Ingress] = new(ingressRegistry)
|
||||
routerReg reg.Registry[router.Router] = new(routerRegistry)
|
||||
@@ -200,6 +202,11 @@ func RateLimiterRegistry() reg.Registry[rate.RateLimiter] {
|
||||
return rateLimiterReg
|
||||
}
|
||||
|
||||
// QuotaLimiterRegistry returns the global registry of quota limiter instances.
|
||||
func QuotaLimiterRegistry() reg.Registry[*quota.Limiter] {
|
||||
return quotaLimiterReg
|
||||
}
|
||||
|
||||
// IngressRegistry returns the global registry of ingress instances.
|
||||
func IngressRegistry() reg.Registry[ingress.Ingress] {
|
||||
return ingressReg
|
||||
|
||||
Reference in New Issue
Block a user