mirror of
https://github.com/go-i2p/go-limit.git
synced 2025-07-13 11:54:45 -04:00
Implement connection limiter
This commit is contained in:
22
LICENSE.md
Normal file
22
LICENSE.md
Normal file
@ -0,0 +1,22 @@
|
|||||||
|
The MIT License (MIT)
|
||||||
|
|
||||||
|
Copyright (c) 2025 idk
|
||||||
|
|
||||||
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
|
furnished to do so, subject to the following conditions:
|
||||||
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
|
copies or substantial portions of the Software.
|
||||||
|
|
||||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
||||||
|
|
2
Makefile
Normal file
2
Makefile
Normal file
@ -0,0 +1,2 @@
|
|||||||
|
fmt:
|
||||||
|
find . -name '*.go' -exec gofumpt -s -w -extra {} \;
|
67
README.md
Normal file
67
README.md
Normal file
@ -0,0 +1,67 @@
|
|||||||
|
# go-limit
|
||||||
|
|
||||||
|
A simple, thread-safe connection limiter for Go's `net.Listener` that manages concurrent connections and rate limiting.
|
||||||
|
|
||||||
|
## Install
|
||||||
|
|
||||||
|
```bash
|
||||||
|
go get github.com/go-i2p/go-limit
|
||||||
|
```
|
||||||
|
|
||||||
|
## Quick Start
|
||||||
|
|
||||||
|
```go
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"log"
|
||||||
|
"net"
|
||||||
|
"github.com/go-i2p/go-limit"
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
base, _ := net.Listen("tcp", ":8080")
|
||||||
|
|
||||||
|
listener := limitedlistener.NewLimitedListener(base,
|
||||||
|
limitedlistener.WithMaxConnections(1000), // max concurrent
|
||||||
|
limitedlistener.WithRateLimit(100)) // per second
|
||||||
|
defer listener.Close()
|
||||||
|
|
||||||
|
for {
|
||||||
|
conn, err := listener.Accept()
|
||||||
|
if err != nil {
|
||||||
|
log.Printf("Accept error: %v", err)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
go handleConnection(conn)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
- Limit concurrent connections
|
||||||
|
- Rate limiting (connections per second)
|
||||||
|
- Connection tracking
|
||||||
|
- Real-time statistics
|
||||||
|
|
||||||
|
## Configuration
|
||||||
|
|
||||||
|
```go
|
||||||
|
// Set maximum concurrent connections
|
||||||
|
WithMaxConnections(1000)
|
||||||
|
|
||||||
|
// Set rate limit (connections per second)
|
||||||
|
WithRateLimit(100.0)
|
||||||
|
|
||||||
|
// Get current stats
|
||||||
|
stats := listener.GetStats()
|
||||||
|
```
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
MIT License
|
||||||
|
|
||||||
|
## Support
|
||||||
|
|
||||||
|
Issues and PRs welcome at [github.com/go-i2p/go-limit](https://github.com/go-i2p/go-limit)
|
5
go.mod
Normal file
5
go.mod
Normal file
@ -0,0 +1,5 @@
|
|||||||
|
module github.com/go-i2p/go-limit
|
||||||
|
|
||||||
|
go 1.23.5
|
||||||
|
|
||||||
|
require golang.org/x/time v0.9.0
|
2
go.sum
Normal file
2
go.sum
Normal file
@ -0,0 +1,2 @@
|
|||||||
|
golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY=
|
||||||
|
golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
|
29
limitedListener.go
Normal file
29
limitedListener.go
Normal file
@ -0,0 +1,29 @@
|
|||||||
|
package limitedlistener
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"golang.org/x/time/rate"
|
||||||
|
)
|
||||||
|
|
||||||
|
// LimitedListener wraps a net.Listener with connection limiting capabilities
|
||||||
|
type LimitedListener struct {
|
||||||
|
net.Listener
|
||||||
|
maxConns int
|
||||||
|
activeConns int64
|
||||||
|
limiter *rate.Limiter
|
||||||
|
mu sync.Mutex
|
||||||
|
done chan struct{}
|
||||||
|
activeSet map[net.Conn]struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Option defines the functional option pattern for configuration
|
||||||
|
type Option func(*LimitedListener)
|
||||||
|
|
||||||
|
// ErrMaxConnsReached is returned when connection limit is reached
|
||||||
|
var ErrMaxConnsReached = fmt.Errorf("maximum connections reached")
|
||||||
|
|
||||||
|
// ErrRateLimitExceeded is returned when rate limit is exceeded
|
||||||
|
var ErrRateLimitExceeded = fmt.Errorf("rate limit exceeded")
|
58
listener.go
Normal file
58
listener.go
Normal file
@ -0,0 +1,58 @@
|
|||||||
|
package limitedlistener
|
||||||
|
|
||||||
|
import "net"
|
||||||
|
|
||||||
|
// NewLimitedListener creates a new connection-limited listener
|
||||||
|
func NewLimitedListener(listener net.Listener, opts ...Option) *LimitedListener {
|
||||||
|
l := &LimitedListener{
|
||||||
|
Listener: listener,
|
||||||
|
maxConns: 1000, // default limit
|
||||||
|
done: make(chan struct{}),
|
||||||
|
activeSet: make(map[net.Conn]struct{}),
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, opt := range opts {
|
||||||
|
opt(l)
|
||||||
|
}
|
||||||
|
|
||||||
|
return l
|
||||||
|
}
|
||||||
|
|
||||||
|
// Accept accepts a connection with limiting logic
|
||||||
|
func (l *LimitedListener) Accept() (net.Conn, error) {
|
||||||
|
for {
|
||||||
|
// Check if rate limit is exceeded
|
||||||
|
if l.limiter != nil {
|
||||||
|
if !l.limiter.Allow() {
|
||||||
|
return nil, ErrRateLimitExceeded
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check concurrent connection limit
|
||||||
|
l.mu.Lock()
|
||||||
|
if l.maxConns > 0 && l.activeConns >= int64(l.maxConns) {
|
||||||
|
l.mu.Unlock()
|
||||||
|
return nil, ErrMaxConnsReached
|
||||||
|
}
|
||||||
|
l.mu.Unlock()
|
||||||
|
|
||||||
|
// Accept the connection
|
||||||
|
conn, err := l.Listener.Accept()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wrap the connection for tracking
|
||||||
|
tracked := &trackedConn{
|
||||||
|
Conn: conn,
|
||||||
|
listener: l,
|
||||||
|
}
|
||||||
|
|
||||||
|
l.mu.Lock()
|
||||||
|
l.activeConns++
|
||||||
|
l.activeSet[tracked] = struct{}{}
|
||||||
|
l.mu.Unlock()
|
||||||
|
|
||||||
|
return tracked, nil
|
||||||
|
}
|
||||||
|
}
|
17
options.go
Normal file
17
options.go
Normal file
@ -0,0 +1,17 @@
|
|||||||
|
package limitedlistener
|
||||||
|
|
||||||
|
import "golang.org/x/time/rate"
|
||||||
|
|
||||||
|
// WithMaxConnections sets the maximum number of concurrent connections
|
||||||
|
func WithMaxConnections(max int) Option {
|
||||||
|
return func(l *LimitedListener) {
|
||||||
|
l.maxConns = max
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithRateLimit sets connections per second limit
|
||||||
|
func WithRateLimit(perSecond float64) Option {
|
||||||
|
return func(l *LimitedListener) {
|
||||||
|
l.limiter = rate.NewLimiter(rate.Limit(perSecond), int(perSecond))
|
||||||
|
}
|
||||||
|
}
|
39
stats.go
Normal file
39
stats.go
Normal file
@ -0,0 +1,39 @@
|
|||||||
|
package limitedlistener
|
||||||
|
|
||||||
|
// Stats provides current listener statistics
|
||||||
|
type Stats struct {
|
||||||
|
ActiveConnections int64
|
||||||
|
MaxConnections int
|
||||||
|
RateLimit float64
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetStats returns current listener statistics
|
||||||
|
func (l *LimitedListener) GetStats() Stats {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
|
||||||
|
var rateLimit float64
|
||||||
|
if l.limiter != nil {
|
||||||
|
rateLimit = float64(l.limiter.Limit())
|
||||||
|
}
|
||||||
|
|
||||||
|
return Stats{
|
||||||
|
ActiveConnections: l.activeConns,
|
||||||
|
MaxConnections: l.maxConns,
|
||||||
|
RateLimit: rateLimit,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close implements graceful shutdown
|
||||||
|
func (l *LimitedListener) Close() error {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
|
||||||
|
// Close all active connections
|
||||||
|
for conn := range l.activeSet {
|
||||||
|
conn.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
close(l.done)
|
||||||
|
return l.Listener.Close()
|
||||||
|
}
|
24
trackedConn.go
Normal file
24
trackedConn.go
Normal file
@ -0,0 +1,24 @@
|
|||||||
|
package limitedlistener
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// trackedConn wraps a net.Conn to track its lifecycle
|
||||||
|
type trackedConn struct {
|
||||||
|
net.Conn
|
||||||
|
listener *LimitedListener
|
||||||
|
once sync.Once
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close implements net.Conn Close with connection tracking
|
||||||
|
func (c *trackedConn) Close() error {
|
||||||
|
c.once.Do(func() {
|
||||||
|
c.listener.mu.Lock()
|
||||||
|
delete(c.listener.activeSet, c)
|
||||||
|
c.listener.activeConns--
|
||||||
|
c.listener.mu.Unlock()
|
||||||
|
})
|
||||||
|
return c.Conn.Close()
|
||||||
|
}
|
Reference in New Issue
Block a user