v0.3.4 rate limit and net limit separated, rate limit by size added

This commit is contained in:
2025-07-13 15:09:40 -04:00
parent cc27f5cc1c
commit be5bb9f2bd
15 changed files with 137 additions and 83 deletions

2
go.mod
View File

@ -4,7 +4,7 @@ go 1.24.5
require ( require (
github.com/lixenwraith/config v0.0.0-20250712170030-7d38402e0497 github.com/lixenwraith/config v0.0.0-20250712170030-7d38402e0497
github.com/lixenwraith/log v0.0.0-20250710220042-139c522da9aa github.com/lixenwraith/log v0.0.0-20250713052337-7b17ce48112f
github.com/panjf2000/gnet/v2 v2.9.1 github.com/panjf2000/gnet/v2 v2.9.1
github.com/valyala/fasthttp v1.63.0 github.com/valyala/fasthttp v1.63.0
) )

4
go.sum
View File

@ -8,8 +8,8 @@ github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zt
github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
github.com/lixenwraith/config v0.0.0-20250712170030-7d38402e0497 h1:ixTIdJSd945n/IhMRwGwQVmQnQ1nUr5z1wn31jXq9FU= github.com/lixenwraith/config v0.0.0-20250712170030-7d38402e0497 h1:ixTIdJSd945n/IhMRwGwQVmQnQ1nUr5z1wn31jXq9FU=
github.com/lixenwraith/config v0.0.0-20250712170030-7d38402e0497/go.mod h1:y7kgDrWIFROWJJ6ASM/SPTRRAj27FjRGWh2SDLcdQ68= github.com/lixenwraith/config v0.0.0-20250712170030-7d38402e0497/go.mod h1:y7kgDrWIFROWJJ6ASM/SPTRRAj27FjRGWh2SDLcdQ68=
github.com/lixenwraith/log v0.0.0-20250710220042-139c522da9aa h1:lIo780MTgZXH5jF+Qr24fRDFQTQ3eu3OaaFYLP9ZkR0= github.com/lixenwraith/log v0.0.0-20250713052337-7b17ce48112f h1:GBSik9B/Eaqz1ogSbG6C+Ti199k8cd7atnNSkCdkfV4=
github.com/lixenwraith/log v0.0.0-20250710220042-139c522da9aa/go.mod h1:lEjTIMWGW+XGn20x5Ec0qkNK4LCd7Y/04PII51crYKk= github.com/lixenwraith/log v0.0.0-20250713052337-7b17ce48112f/go.mod h1:lEjTIMWGW+XGn20x5Ec0qkNK4LCd7Y/04PII51crYKk=
github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
github.com/panjf2000/ants/v2 v2.11.3 h1:AfI0ngBoXJmYOpDh9m516vjqoUu2sLrIVgppI9TZVpg= github.com/panjf2000/ants/v2 v2.11.3 h1:AfI0ngBoXJmYOpDh9m516vjqoUu2sLrIVgppI9TZVpg=

View File

@ -37,10 +37,6 @@ type SourceConfig struct {
// Type-specific configuration options // Type-specific configuration options
Options map[string]any `toml:"options"` Options map[string]any `toml:"options"`
// Placeholder for future source-side rate limiting
// This will be used for features like aggregation and summarization
NetLimit *NetLimitConfig `toml:"net_limit"`
} }
// SinkConfig represents an output destination // SinkConfig represents an output destination
@ -116,6 +112,13 @@ func validateSource(pipelineName string, sourceIndex int, cfg *SourceConfig) err
} }
} }
// Validate net_limit if present within Options
if rl, ok := cfg.Options["net_limit"].(map[string]any); ok {
if err := validateNetLimitOptions("HTTP source", pipelineName, sourceIndex, rl); err != nil {
return err
}
}
case "tcp": case "tcp":
// Validate TCP source options // Validate TCP source options
port, ok := toInt(cfg.Options["port"]) port, ok := toInt(cfg.Options["port"])
@ -124,6 +127,13 @@ func validateSource(pipelineName string, sourceIndex int, cfg *SourceConfig) err
pipelineName, sourceIndex) pipelineName, sourceIndex)
} }
// Validate net_limit if present within Options
if rl, ok := cfg.Options["net_limit"].(map[string]any); ok {
if err := validateNetLimitOptions("TCP source", pipelineName, sourceIndex, rl); err != nil {
return err
}
}
default: default:
return fmt.Errorf("pipeline '%s' source[%d]: unknown source type '%s'", return fmt.Errorf("pipeline '%s' source[%d]: unknown source type '%s'",
pipelineName, sourceIndex, cfg.Type) pipelineName, sourceIndex, cfg.Type)

View File

@ -24,6 +24,8 @@ type RateLimitConfig struct {
Burst float64 `toml:"burst"` Burst float64 `toml:"burst"`
// Policy defines the action to take when the limit is exceeded. "pass" or "drop". // Policy defines the action to take when the limit is exceeded. "pass" or "drop".
Policy string `toml:"policy"` Policy string `toml:"policy"`
// MaxEntrySizeBytes is the maximum allowed size for a single log entry. 0 = no limit.
MaxEntrySizeBytes int `toml:"max_entry_size_bytes"`
} }
func validateRateLimit(pipelineName string, cfg *RateLimitConfig) error { func validateRateLimit(pipelineName string, cfg *RateLimitConfig) error {
@ -39,6 +41,10 @@ func validateRateLimit(pipelineName string, cfg *RateLimitConfig) error {
return fmt.Errorf("pipeline '%s': rate limit burst cannot be negative", pipelineName) return fmt.Errorf("pipeline '%s': rate limit burst cannot be negative", pipelineName)
} }
if cfg.MaxEntrySizeBytes < 0 {
return fmt.Errorf("pipeline '%s': max entry size bytes cannot be negative", pipelineName)
}
// Validate policy // Validate policy
switch strings.ToLower(cfg.Policy) { switch strings.ToLower(cfg.Policy) {
case "", "pass", "drop": case "", "pass", "drop":

View File

@ -15,7 +15,7 @@ type TCPConfig struct {
NetLimit *NetLimitConfig `toml:"net_limit"` NetLimit *NetLimitConfig `toml:"net_limit"`
// Heartbeat // Heartbeat
Heartbeat HeartbeatConfig `toml:"heartbeat"` Heartbeat *HeartbeatConfig `toml:"heartbeat"`
} }
type HTTPConfig struct { type HTTPConfig struct {
@ -34,7 +34,7 @@ type HTTPConfig struct {
NetLimit *NetLimitConfig `toml:"net_limit"` NetLimit *NetLimitConfig `toml:"net_limit"`
// Heartbeat // Heartbeat
Heartbeat HeartbeatConfig `toml:"heartbeat"` Heartbeat *HeartbeatConfig `toml:"heartbeat"`
} }
type HeartbeatConfig struct { type HeartbeatConfig struct {

View File

@ -1,16 +1,17 @@
// FILE: src/internal/netlimit/netlimiter.go // FILE: src/internal/limiter/token_bucket.go
package netlimit package limiter
import ( import (
"sync" "sync"
"time" "time"
) )
// TokenBucket implements a token bucket net limiter // TokenBucket implements a token bucket rate limiter
// Safe for concurrent use.
type TokenBucket struct { type TokenBucket struct {
capacity float64 capacity float64
tokens float64 tokens float64
refillRate float64 refillRate float64 // tokens per second
lastRefill time.Time lastRefill time.Time
mu sync.Mutex mu sync.Mutex
} }
@ -19,7 +20,7 @@ type TokenBucket struct {
func NewTokenBucket(capacity float64, refillRate float64) *TokenBucket { func NewTokenBucket(capacity float64, refillRate float64) *TokenBucket {
return &TokenBucket{ return &TokenBucket{
capacity: capacity, capacity: capacity,
tokens: capacity, tokens: capacity, // Start full
refillRate: refillRate, refillRate: refillRate,
lastRefill: time.Now(), lastRefill: time.Now(),
} }
@ -35,7 +36,27 @@ func (tb *TokenBucket) AllowN(n float64) bool {
tb.mu.Lock() tb.mu.Lock()
defer tb.mu.Unlock() defer tb.mu.Unlock()
// Refill tokens based on time elapsed tb.refill()
if tb.tokens >= n {
tb.tokens -= n
return true
}
return false
}
// Tokens returns the current number of available tokens
func (tb *TokenBucket) Tokens() float64 {
tb.mu.Lock()
defer tb.mu.Unlock()
tb.refill()
return tb.tokens
}
// refill adds tokens based on time elapsed since last refill
// MUST be called with mutex held
func (tb *TokenBucket) refill() {
now := time.Now() now := time.Now()
elapsed := now.Sub(tb.lastRefill).Seconds() elapsed := now.Sub(tb.lastRefill).Seconds()
@ -43,7 +64,6 @@ func (tb *TokenBucket) AllowN(n float64) bool {
if elapsed < 0 { if elapsed < 0 {
// Clock went backwards, reset to current time but don't add tokens // Clock went backwards, reset to current time but don't add tokens
tb.lastRefill = now tb.lastRefill = now
// Don't log here as this is a hot path
elapsed = 0 elapsed = 0
} }
@ -52,11 +72,4 @@ func (tb *TokenBucket) AllowN(n float64) bool {
tb.tokens = tb.capacity tb.tokens = tb.capacity
} }
tb.lastRefill = now tb.lastRefill = now
// Check if we have enough tokens
if tb.tokens >= n {
tb.tokens -= n
return true
}
return false
} }

View File

@ -10,11 +10,12 @@ import (
"time" "time"
"logwisp/src/internal/config" "logwisp/src/internal/config"
"logwisp/src/internal/limiter"
"github.com/lixenwraith/log" "github.com/lixenwraith/log"
) )
// Manages net limiting for a transport // Limiter manages net limiting for a transport
type Limiter struct { type Limiter struct {
config config.NetLimitConfig config config.NetLimitConfig
logger *log.Logger logger *log.Logger
@ -24,7 +25,7 @@ type Limiter struct {
ipMu sync.RWMutex ipMu sync.RWMutex
// Global limiter for the transport // Global limiter for the transport
globalLimiter *TokenBucket globalLimiter *limiter.TokenBucket
// Connection tracking // Connection tracking
ipConnections map[string]*atomic.Int32 ipConnections map[string]*atomic.Int32
@ -46,7 +47,7 @@ type Limiter struct {
} }
type ipLimiter struct { type ipLimiter struct {
bucket *TokenBucket bucket *limiter.TokenBucket
lastSeen time.Time lastSeen time.Time
connections atomic.Int32 connections atomic.Int32
} }
@ -76,7 +77,7 @@ func New(cfg config.NetLimitConfig, logger *log.Logger) *Limiter {
// Create global limiter if not using per-IP limiting // Create global limiter if not using per-IP limiting
if cfg.LimitBy == "global" { if cfg.LimitBy == "global" {
l.globalLimiter = NewTokenBucket( l.globalLimiter = limiter.NewTokenBucket(
float64(cfg.BurstSize), float64(cfg.BurstSize),
cfg.RequestsPerSecond, cfg.RequestsPerSecond,
) )
@ -334,24 +335,24 @@ func (l *Limiter) checkLimit(ip string) bool {
case "ip", "": case "ip", "":
// Default to per-IP limiting // Default to per-IP limiting
l.ipMu.Lock() l.ipMu.Lock()
limiter, exists := l.ipLimiters[ip] lim, exists := l.ipLimiters[ip]
if !exists { if !exists {
// Create new limiter for this IP // Create new limiter for this IP
limiter = &ipLimiter{ lim = &ipLimiter{
bucket: NewTokenBucket( bucket: limiter.NewTokenBucket(
float64(l.config.BurstSize), float64(l.config.BurstSize),
l.config.RequestsPerSecond, l.config.RequestsPerSecond,
), ),
lastSeen: time.Now(), lastSeen: time.Now(),
} }
l.ipLimiters[ip] = limiter l.ipLimiters[ip] = lim
l.uniqueIPs.Add(1) l.uniqueIPs.Add(1)
l.logger.Debug("msg", "Created new IP limiter", l.logger.Debug("msg", "Created new IP limiter",
"ip", ip, "ip", ip,
"total_ips", l.uniqueIPs.Load()) "total_ips", l.uniqueIPs.Load())
} else { } else {
limiter.lastSeen = time.Now() lim.lastSeen = time.Now()
} }
l.ipMu.Unlock() l.ipMu.Unlock()
@ -366,7 +367,7 @@ func (l *Limiter) checkLimit(ip string) bool {
} }
} }
return limiter.bucket.Allow() return lim.bucket.Allow()
default: default:
// Unknown limit_by value, allow by default // Unknown limit_by value, allow by default
@ -398,8 +399,8 @@ func (l *Limiter) cleanup() {
defer l.ipMu.Unlock() defer l.ipMu.Unlock()
cleaned := 0 cleaned := 0
for ip, limiter := range l.ipLimiters { for ip, lim := range l.ipLimiters {
if now.Sub(limiter.lastSeen) > staleTimeout { if now.Sub(lim.lastSeen) > staleTimeout {
delete(l.ipLimiters, ip) delete(l.ipLimiters, ip)
cleaned++ cleaned++
} }

View File

@ -3,27 +3,25 @@ package ratelimit
import ( import (
"strings" "strings"
"sync"
"sync/atomic" "sync/atomic"
"time"
"logwisp/src/internal/config"
"logwisp/src/internal/limiter"
"logwisp/src/internal/source"
"github.com/lixenwraith/log" "github.com/lixenwraith/log"
"logwisp/src/internal/config"
"logwisp/src/internal/source"
) )
// Limiter enforces rate limits on log entries flowing through a pipeline. // Limiter enforces rate limits on log entries flowing through a pipeline.
type Limiter struct { type Limiter struct {
mu sync.Mutex bucket *limiter.TokenBucket
rate float64 policy config.RateLimitPolicy
burst float64 logger *log.Logger
tokens float64
lastToken time.Time
policy config.RateLimitPolicy
logger *log.Logger
// Statistics // Statistics
droppedCount atomic.Uint64 maxEntrySizeBytes int
droppedBySizeCount atomic.Uint64
droppedCount atomic.Uint64
} }
// New creates a new rate limiter. If cfg.Rate is 0, it returns nil. // New creates a new rate limiter. If cfg.Rate is 0, it returns nil.
@ -46,12 +44,14 @@ func New(cfg config.RateLimitConfig, logger *log.Logger) (*Limiter, error) {
} }
l := &Limiter{ l := &Limiter{
rate: cfg.Rate, bucket: limiter.NewTokenBucket(burst, cfg.Rate),
burst: burst, policy: policy,
tokens: burst, logger: logger,
lastToken: time.Now(), maxEntrySizeBytes: cfg.MaxEntrySizeBytes,
policy: policy, }
logger: logger,
if cfg.Rate > 0 {
l.bucket = limiter.NewTokenBucket(burst, cfg.Rate)
} }
return l, nil return l, nil
@ -60,46 +60,51 @@ func New(cfg config.RateLimitConfig, logger *log.Logger) (*Limiter, error) {
// Allow checks if a log entry is allowed to pass based on the rate limit. // Allow checks if a log entry is allowed to pass based on the rate limit.
// It returns true if the entry should pass, false if it should be dropped. // It returns true if the entry should pass, false if it should be dropped.
func (l *Limiter) Allow(entry source.LogEntry) bool { func (l *Limiter) Allow(entry source.LogEntry) bool {
if l.policy == config.PolicyPass { if l == nil || l.policy == config.PolicyPass {
return true return true
} }
l.mu.Lock() // Check size limit first
defer l.mu.Unlock() if l.maxEntrySizeBytes > 0 && entry.RawSize > l.maxEntrySizeBytes {
l.droppedBySizeCount.Add(1)
now := time.Now() return false
elapsed := now.Sub(l.lastToken).Seconds()
if elapsed < 0 {
// Clock went backwards, don't add tokens
l.lastToken = now
elapsed = 0
} }
l.tokens += elapsed * l.rate // Check rate limit if configured
if l.tokens > l.burst { if l.bucket != nil {
l.tokens = l.burst if l.bucket.Allow() {
} return true
l.lastToken = now }
// Not enough tokens, drop the entry
if l.tokens >= 1 { l.droppedCount.Add(1)
l.tokens-- return false
return true
} }
// Not enough tokens, drop the entry // No rate limit configured, size check passed
l.droppedCount.Add(1) return true
return false
} }
// GetStats returns the statistics for the limiter. // GetStats returns the statistics for the limiter.
func (l *Limiter) GetStats() map[string]any { func (l *Limiter) GetStats() map[string]any {
return map[string]any{ if l == nil {
"dropped_total": l.droppedCount.Load(), return map[string]any{
"policy": policyString(l.policy), "enabled": false,
"rate": l.rate, }
"burst": l.burst,
} }
stats := map[string]any{
"enabled": true,
"dropped_total": l.droppedCount.Load(),
"dropped_by_size_total": l.droppedBySizeCount.Load(),
"policy": policyString(l.policy),
"max_entry_size_bytes": l.maxEntrySizeBytes,
}
if l.bucket != nil {
stats["tokens"] = l.bucket.Tokens()
}
return stats
} }
// policyString returns the string representation of the policy. // policyString returns the string representation of the policy.

View File

@ -54,7 +54,7 @@ type HTTPConfig struct {
BufferSize int BufferSize int
StreamPath string StreamPath string
StatusPath string StatusPath string
Heartbeat config.HeartbeatConfig Heartbeat *config.HeartbeatConfig
SSL *config.SSLConfig SSL *config.SSLConfig
NetLimit *config.NetLimitConfig NetLimit *config.NetLimitConfig
} }
@ -84,6 +84,7 @@ func NewHTTPSink(options map[string]any, logger *log.Logger) (*HTTPSink, error)
// Extract heartbeat config // Extract heartbeat config
if hb, ok := options["heartbeat"].(map[string]any); ok { if hb, ok := options["heartbeat"].(map[string]any); ok {
cfg.Heartbeat = &config.HeartbeatConfig{}
cfg.Heartbeat.Enabled, _ = hb["enabled"].(bool) cfg.Heartbeat.Enabled, _ = hb["enabled"].(bool)
if interval, ok := toInt(hb["interval_seconds"]); ok { if interval, ok := toInt(hb["interval_seconds"]); ok {
cfg.Heartbeat.IntervalSeconds = interval cfg.Heartbeat.IntervalSeconds = interval

View File

@ -41,7 +41,7 @@ type TCPSink struct {
type TCPConfig struct { type TCPConfig struct {
Port int Port int
BufferSize int BufferSize int
Heartbeat config.HeartbeatConfig Heartbeat *config.HeartbeatConfig
SSL *config.SSLConfig SSL *config.SSLConfig
NetLimit *config.NetLimitConfig NetLimit *config.NetLimitConfig
} }
@ -63,6 +63,7 @@ func NewTCPSink(options map[string]any, logger *log.Logger) (*TCPSink, error) {
// Extract heartbeat config // Extract heartbeat config
if hb, ok := options["heartbeat"].(map[string]any); ok { if hb, ok := options["heartbeat"].(map[string]any); ok {
cfg.Heartbeat = &config.HeartbeatConfig{}
cfg.Heartbeat.Enabled, _ = hb["enabled"].(bool) cfg.Heartbeat.Enabled, _ = hb["enabled"].(bool)
if interval, ok := toInt(hb["interval_seconds"]); ok { if interval, ok := toInt(hb["interval_seconds"]); ok {
cfg.Heartbeat.IntervalSeconds = interval cfg.Heartbeat.IntervalSeconds = interval

View File

@ -258,7 +258,10 @@ func (w *fileWatcher) checkFile() error {
continue continue
} }
rawSize := len(line)
entry := w.parseLine(line) entry := w.parseLine(line)
entry.RawSize = rawSize
w.callback(entry) w.callback(entry)
w.entriesRead.Add(1) w.entriesRead.Add(1)
w.lastReadTime.Store(time.Now()) w.lastReadTime.Store(time.Now())

View File

@ -271,6 +271,7 @@ func (h *HTTPSource) parseEntries(body []byte) ([]LogEntry, error) {
if single.Source == "" { if single.Source == "" {
single.Source = "http" single.Source = "http"
} }
single.RawSize = len(body)
entries = append(entries, single) entries = append(entries, single)
return entries, nil return entries, nil
} }
@ -278,6 +279,8 @@ func (h *HTTPSource) parseEntries(body []byte) ([]LogEntry, error) {
// Try to parse as JSON array // Try to parse as JSON array
var array []LogEntry var array []LogEntry
if err := json.Unmarshal(body, &array); err == nil { if err := json.Unmarshal(body, &array); err == nil {
// TODO: Placeholder; For array, divide total size by entry count as approximation
approxSizePerEntry := len(body) / len(array)
for i, entry := range array { for i, entry := range array {
if entry.Message == "" { if entry.Message == "" {
return nil, fmt.Errorf("entry %d missing required field: message", i) return nil, fmt.Errorf("entry %d missing required field: message", i)
@ -288,6 +291,8 @@ func (h *HTTPSource) parseEntries(body []byte) ([]LogEntry, error) {
if entry.Source == "" { if entry.Source == "" {
array[i].Source = "http" array[i].Source = "http"
} }
// TODO: Placeholder
array[i].RawSize = approxSizePerEntry
} }
return array, nil return array, nil
} }
@ -313,6 +318,7 @@ func (h *HTTPSource) parseEntries(body []byte) ([]LogEntry, error) {
if entry.Source == "" { if entry.Source == "" {
entry.Source = "http" entry.Source = "http"
} }
entry.RawSize = len(line)
entries = append(entries, entry) entries = append(entries, entry)
} }

View File

@ -13,6 +13,7 @@ type LogEntry struct {
Level string `json:"level,omitempty"` Level string `json:"level,omitempty"`
Message string `json:"message"` Message string `json:"message"`
Fields json.RawMessage `json:"fields,omitempty"` Fields json.RawMessage `json:"fields,omitempty"`
RawSize int `json:"-"`
} }
// Source represents an input data stream // Source represents an input data stream

View File

@ -82,6 +82,7 @@ func (s *StdinSource) readLoop() {
Source: "stdin", Source: "stdin",
Message: line, Message: line,
Level: extractLogLevel(line), Level: extractLogLevel(line),
RawSize: len(line),
} }
s.publish(entry) s.publish(entry)

View File

@ -341,6 +341,9 @@ func (s *tcpSourceServer) OnTraffic(c gnet.Conn) gnet.Action {
continue continue
} }
// Capture raw line size before parsing
rawSize := len(line)
// Parse JSON log entry // Parse JSON log entry
var entry LogEntry var entry LogEntry
if err := json.Unmarshal(line, &entry); err != nil { if err := json.Unmarshal(line, &entry); err != nil {
@ -364,6 +367,9 @@ func (s *tcpSourceServer) OnTraffic(c gnet.Conn) gnet.Action {
entry.Source = "tcp" entry.Source = "tcp"
} }
// Set raw size
entry.RawSize = rawSize
// Publish the entry // Publish the entry
s.source.publish(entry) s.source.publish(entry)
} }