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

limiter/quota/wrapper/conn.go

@@ -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 &quotaConn{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
+}

limiter/quota/wrapper/listener.go

@@ -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 &quotaListener{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()
+}