144 lines
3.4 KiB
Go
144 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"log/slog"
|
|
"os"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RequestCounter tracks request counts per path in memory with periodic snapshots
|
|
type RequestCounter struct {
|
|
counts map[string]int64
|
|
mu sync.RWMutex
|
|
snapshotPath string
|
|
stopChan chan struct{}
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// NewRequestCounter creates a new request counter with periodic snapshots
|
|
func NewRequestCounter(snapshotInterval time.Duration) (*RequestCounter, error) {
|
|
snapshotPath := os.Getenv("REQUEST_COUNTS_PATH")
|
|
if snapshotPath == "" {
|
|
err := errors.New("REQUEST_COUNTS_PATH environment variable must be set and non-empty")
|
|
slog.Error("failed to initialize request counter", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
c := &RequestCounter{
|
|
counts: make(map[string]int64),
|
|
snapshotPath: snapshotPath,
|
|
stopChan: make(chan struct{}),
|
|
}
|
|
|
|
// Load existing counts from snapshot file
|
|
if err := c.load(); err != nil {
|
|
slog.Warn("failed to load counter snapshot, starting fresh", "error", err, "path", snapshotPath)
|
|
}
|
|
|
|
// Start periodic snapshot goroutine
|
|
c.wg.Add(1)
|
|
go c.periodicSnapshot(snapshotInterval)
|
|
|
|
return c, nil
|
|
}
|
|
|
|
// Increment increments the counter for a given path and returns the new count
|
|
func (c *RequestCounter) Increment(path string) int64 {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.counts[path]++
|
|
return c.counts[path]
|
|
}
|
|
|
|
// Get returns the current count for a given path
|
|
func (c *RequestCounter) Get(path string) int64 {
|
|
c.mu.RLock()
|
|
defer c.mu.RUnlock()
|
|
return c.counts[path]
|
|
}
|
|
|
|
func (c *RequestCounter) GetTotal() int64 {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
var total int64
|
|
for _, count := range c.counts {
|
|
total += count
|
|
}
|
|
return total
|
|
}
|
|
|
|
// load reads the snapshot file and loads counts into memory
|
|
func (c *RequestCounter) load() error {
|
|
data, err := os.ReadFile(c.snapshotPath)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return nil // File doesn't exist yet, that's okay
|
|
}
|
|
return err
|
|
}
|
|
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
if err := json.Unmarshal(data, &c.counts); err != nil {
|
|
return err
|
|
}
|
|
|
|
slog.Info("loaded request counter snapshot", "paths", len(c.counts), "file", c.snapshotPath)
|
|
return nil
|
|
}
|
|
|
|
// snapshot writes current counts to disk
|
|
func (c *RequestCounter) snapshot() error {
|
|
c.mu.RLock()
|
|
data, err := json.MarshalIndent(c.counts, "", " ")
|
|
c.mu.RUnlock()
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Write to temp file first, then rename for atomicity
|
|
tempPath := c.snapshotPath + ".tmp"
|
|
if err := os.WriteFile(tempPath, data, 0644); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := os.Rename(tempPath, c.snapshotPath); err != nil {
|
|
return err
|
|
}
|
|
|
|
slog.Debug("saved request counter snapshot", "paths", len(c.counts), "file", c.snapshotPath)
|
|
return nil
|
|
}
|
|
|
|
// periodicSnapshot runs in a goroutine and saves snapshots at regular intervals
|
|
func (c *RequestCounter) periodicSnapshot(interval time.Duration) {
|
|
defer c.wg.Done()
|
|
ticker := time.NewTicker(interval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
if err := c.snapshot(); err != nil {
|
|
slog.Error("failed to save counter snapshot", "error", err)
|
|
}
|
|
case <-c.stopChan:
|
|
// Final snapshot before shutdown
|
|
if err := c.snapshot(); err != nil {
|
|
slog.Error("failed to save final counter snapshot", "error", err)
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// Close stops the periodic snapshot goroutine and saves a final snapshot
|
|
func (c *RequestCounter) Close() error {
|
|
close(c.stopChan)
|
|
c.wg.Wait()
|
|
return nil
|
|
}
|