devplacepy/nodes/client/limits.go

package main

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	mebibyte        = 1024 * 1024
	persistInterval = 5 * time.Second
)

type usageState struct {
	Day   string `json:"day"`
	Bytes int64  `json:"bytes"`
}

type Limiter struct {
	mu          sync.Mutex
	cfg         Config
	loc         *time.Location
	now         func() time.Time
	statePath   string
	serverLimit ServerLimits
	serverSet   ServerSettings
	hasServer   bool
	day         string
	used        int64
	active      int
	lastPersist time.Time
}

func NewLimiter(cfg Config, statePath string, now func() time.Time) (*Limiter, error) {
	loc, err := cfg.Location()
	if err != nil {
		return nil, err
	}
	limiter := &Limiter{cfg: cfg, loc: loc, now: now, statePath: statePath}
	limiter.load()
	return limiter, nil
}

func (l *Limiter) load() {
	if l.statePath == "" {
		return
	}
	data, err := os.ReadFile(l.statePath)
	if err != nil || len(data) > 4096 {
		return
	}
	var state usageState
	if json.Unmarshal(data, &state) != nil || state.Bytes < 0 {
		return
	}
	l.day, l.used = state.Day, state.Bytes
}

func (l *Limiter) persist(force bool) {
	if l.statePath == "" {
		return
	}
	moment := l.now()
	if !force && moment.Sub(l.lastPersist) < persistInterval {
		return
	}
	l.lastPersist = moment
	data, err := json.Marshal(usageState{Day: l.day, Bytes: l.used})
	if err != nil {
		return
	}
	if err := os.MkdirAll(filepath.Dir(l.statePath), 0o700); err != nil {
		return
	}
	_ = writeFileAtomic(l.statePath, data)
}

func (l *Limiter) Flush() {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.roll()
	l.persist(true)
}

func (l *Limiter) ApplyServer(limits *ServerLimits, settings *ServerSettings) {
	l.mu.Lock()
	defer l.mu.Unlock()
	if limits != nil {
		l.serverLimit = *limits
	}
	if settings != nil {
		l.serverSet = *settings
	}
	l.hasServer = limits != nil || settings != nil
}

func (l *Limiter) roll() {
	today := l.now().In(l.loc).Format("2006-01-02")
	if l.day != today {
		l.day = today
		l.used = 0
	}
}

func minPositive(values ...int) int {
	best := 0
	for _, value := range values {
		if value > 0 && (best == 0 || value < best) {
			best = value
		}
	}
	return best
}

func (l *Limiter) capBytes() int64 {
	mb := minPositive(l.cfg.BandwidthMBDay, l.serverSet.BandwidthMBDay, l.serverLimit.MaxDonorMBPerDay)
	return int64(mb) * mebibyte
}

func (l *Limiter) concurrencyCap() int {
	return minPositive(l.cfg.MaxConcurrent, l.serverSet.MaxConcurrent, l.serverLimit.MaxConcurrentPerNode)
}

func minutesOf(text string) int {
	var hours, minutes int
	if len(text) != 5 {
		return -1
	}
	hours = int(text[0]-'0')*10 + int(text[1]-'0')
	minutes = int(text[3]-'0')*10 + int(text[4]-'0')
	return hours*60 + minutes
}

func InWindow(from, to string, moment time.Time) bool {
	if from == "" || to == "" || from == to {
		return true
	}
	start, end := minutesOf(from), minutesOf(to)
	if start < 0 || end < 0 {
		return true
	}
	now := moment.Hour()*60 + moment.Minute()
	if start < end {
		return now >= start && now < end
	}
	return now >= start || now < end
}

func (l *Limiter) hoursCode() string {
	moment := l.now()
	if !InWindow(l.cfg.ActiveFrom, l.cfg.ActiveTo, moment.In(l.loc)) {
		return RefuseHours
	}
	if l.serverSet.ActiveFrom != "" {
		zone, err := time.LoadLocation(l.serverSet.Timezone)
		if err != nil {
			zone = time.UTC
		}
		if !InWindow(l.serverSet.ActiveFrom, l.serverSet.ActiveTo, moment.In(zone)) {
			return RefuseHours
		}
	}
	return ""
}

func (l *Limiter) Acquire() string {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.roll()
	if code := l.hoursCode(); code != "" {
		return code
	}
	if l.used >= l.capBytes() {
		return RefuseBandwidth
	}
	if l.active >= l.concurrencyCap() {
		return RefuseConcurrency
	}
	l.active++
	return ""
}

func (l *Limiter) Release() {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.active > 0 {
		l.active--
	}
}

func (l *Limiter) Remaining() int64 {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.roll()
	left := l.capBytes() - l.used
	if left < 0 {
		return 0
	}
	return left
}

func (l *Limiter) AddBytes(n int64) {
	if n <= 0 {
		return
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	l.roll()
	l.used += n
	l.persist(false)
}

func (l *Limiter) Usage() (day string, used, capacity int64) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.roll()
	return l.day, l.used, l.capBytes()
}

func ReadUsage(path string, loc *time.Location, now time.Time) int64 {
	data, err := os.ReadFile(path)
	if err != nil || len(data) > 4096 {
		return 0
	}
	var state usageState
	if json.Unmarshal(data, &state) != nil || state.Day != now.In(loc).Format("2006-01-02") {
		return 0
	}
	return state.Bytes
}

func (l *Limiter) ConcurrencyCap() int {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.concurrencyCap()
}