feat: add quota limiter and quota API (#107)
This commit is contained in:
Usishchev Yury
GitHub
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user