385 lines
8.8 KiB
Go
385 lines
8.8 KiB
Go
package traffic
|
||
|
||
import (
|
||
"bufio"
|
||
"context"
|
||
"database/sql"
|
||
"io"
|
||
"os"
|
||
"sort"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/havline/havline/internal/logstore"
|
||
"github.com/havline/havline/internal/notify"
|
||
"github.com/havline/havline/internal/proxy"
|
||
)
|
||
|
||
const (
|
||
rateWindow = 5 * time.Second
|
||
activeWindow = 65 * time.Second
|
||
persistInterval = 30 * time.Second
|
||
maxSamplesPerHost = 256
|
||
)
|
||
|
||
type RuleTraffic struct {
|
||
RuleID int64 `json:"rule_id"`
|
||
UploadTotal int64 `json:"upload_total"`
|
||
DownloadTotal int64 `json:"download_total"`
|
||
UploadRate float64 `json:"upload_rate"`
|
||
DownloadRate float64 `json:"download_rate"`
|
||
Connections int `json:"connections"`
|
||
}
|
||
|
||
type ClientConn struct {
|
||
IP string `json:"ip"`
|
||
LastSeen string `json:"last_seen"`
|
||
}
|
||
|
||
type Collector struct {
|
||
path string
|
||
store *Store
|
||
notify *notify.Service
|
||
|
||
mu sync.RWMutex
|
||
hosts map[string]*hostState
|
||
}
|
||
|
||
type hostState struct {
|
||
uploadTotal int64
|
||
downloadTotal int64
|
||
clients map[string]time.Time
|
||
samples []sample
|
||
}
|
||
|
||
type sample struct {
|
||
at time.Time
|
||
up int64
|
||
down int64
|
||
}
|
||
|
||
func NewCollector(db *sql.DB, accessLogPath string, notifySvc *notify.Service) *Collector {
|
||
return &Collector{
|
||
path: accessLogPath,
|
||
store: NewStore(db),
|
||
notify: notifySvc,
|
||
hosts: map[string]*hostState{},
|
||
}
|
||
}
|
||
|
||
func (c *Collector) Start(ctx context.Context) {
|
||
if totals, err := c.store.LoadTotals(ctx); err == nil {
|
||
c.mu.Lock()
|
||
for host, t := range totals {
|
||
st := c.ensureHost(host)
|
||
st.uploadTotal = t.Upload
|
||
st.downloadTotal = t.Download
|
||
}
|
||
c.mu.Unlock()
|
||
}
|
||
|
||
go c.tailLoop(ctx)
|
||
go c.persistLoop(ctx)
|
||
}
|
||
|
||
func (c *Collector) tailLoop(ctx context.Context) {
|
||
offset := c.logOffset()
|
||
for {
|
||
if ctx.Err() != nil {
|
||
return
|
||
}
|
||
offset = c.readFrom(offset)
|
||
time.Sleep(500 * time.Millisecond)
|
||
}
|
||
}
|
||
|
||
func (c *Collector) logOffset() int64 {
|
||
info, err := os.Stat(c.path)
|
||
if err != nil {
|
||
return 0
|
||
}
|
||
return info.Size()
|
||
}
|
||
|
||
func (c *Collector) readFrom(offset int64) int64 {
|
||
file, err := os.Open(c.path)
|
||
if err != nil {
|
||
return offset
|
||
}
|
||
defer file.Close()
|
||
|
||
info, err := file.Stat()
|
||
if err != nil {
|
||
return offset
|
||
}
|
||
if info.Size() < offset {
|
||
offset = 0
|
||
}
|
||
if _, err := file.Seek(offset, 0); err != nil {
|
||
return offset
|
||
}
|
||
|
||
reader := bufio.NewReader(file)
|
||
for {
|
||
line, err := reader.ReadString('\n')
|
||
if len(line) > 0 {
|
||
c.ingestLine(strings.TrimRight(line, "\r\n"), true)
|
||
}
|
||
if err != nil {
|
||
if err == io.EOF {
|
||
pos, seekErr := file.Seek(0, io.SeekCurrent)
|
||
if seekErr == nil {
|
||
return pos
|
||
}
|
||
return info.Size()
|
||
}
|
||
return offset
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Collector) persistLoop(ctx context.Context) {
|
||
ticker := time.NewTicker(persistInterval)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
_ = c.persist(context.Background())
|
||
return
|
||
case <-ticker.C:
|
||
_ = c.persist(ctx)
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Collector) persist(ctx context.Context) error {
|
||
c.mu.RLock()
|
||
snapshot := make(map[string]hostTotals, len(c.hosts))
|
||
for host, st := range c.hosts {
|
||
snapshot[host] = hostTotals{Upload: st.uploadTotal, Download: st.downloadTotal}
|
||
}
|
||
c.mu.RUnlock()
|
||
return c.store.SaveTotals(ctx, snapshot)
|
||
}
|
||
|
||
func (c *Collector) ingestLine(line string, trackLive bool) {
|
||
entry, ok := logstore.ParseAccess(line)
|
||
if !ok || entry.Domain == "" {
|
||
return
|
||
}
|
||
host := EndpointKey(entry.Domain, entry.ServerPort)
|
||
up := entry.RequestLength
|
||
down := entry.BytesSent
|
||
if up == 0 && down == 0 && !trackLive {
|
||
return
|
||
}
|
||
|
||
now := time.Now()
|
||
hitIP := ""
|
||
c.mu.Lock()
|
||
st := c.ensureHost(host)
|
||
st.uploadTotal += up
|
||
st.downloadTotal += down
|
||
if trackLive {
|
||
if up > 0 || down > 0 {
|
||
st.samples = append(st.samples, sample{at: now, up: up, down: down})
|
||
if len(st.samples) > maxSamplesPerHost {
|
||
st.samples = st.samples[len(st.samples)-maxSamplesPerHost:]
|
||
}
|
||
}
|
||
if entry.ClientIP != "" && entry.ClientIP != "-" {
|
||
st.clients[entry.ClientIP] = now
|
||
// 锁内只记 IP：RecordAccessHit 会写库并判断通知，放在锁外，避免阻塞读侧
|
||
hitIP = entry.ClientIP
|
||
}
|
||
c.pruneHost(st, now)
|
||
}
|
||
c.mu.Unlock()
|
||
if hitIP != "" && c.notify != nil {
|
||
c.notify.RecordAccessHit(context.Background(), hitIP)
|
||
}
|
||
}
|
||
|
||
func (c *Collector) ensureHost(host string) *hostState {
|
||
st, ok := c.hosts[host]
|
||
if !ok {
|
||
st = &hostState{clients: map[string]time.Time{}}
|
||
c.hosts[host] = st
|
||
}
|
||
return st
|
||
}
|
||
|
||
func (c *Collector) pruneHost(st *hostState, now time.Time) {
|
||
cutoff := now.Add(-rateWindow)
|
||
for i := 0; i < len(st.samples); {
|
||
if st.samples[i].at.Before(cutoff) {
|
||
st.samples = append(st.samples[:i], st.samples[i+1:]...)
|
||
continue
|
||
}
|
||
i++
|
||
}
|
||
activeCutoff := now.Add(-activeWindow)
|
||
for ip, seen := range st.clients {
|
||
if seen.Before(activeCutoff) {
|
||
delete(st.clients, ip)
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Collector) aggregateEndpoints(endpoints []proxy.Endpoint, now time.Time) (uploadTotal, downloadTotal int64, uploadRate, downloadRate float64, connections int) {
|
||
keys := make([]string, 0, len(endpoints))
|
||
for _, ep := range endpoints {
|
||
keys = append(keys, EndpointKey(ep.Hostname, ep.Port))
|
||
}
|
||
return c.aggregateHostKeys(keys, now)
|
||
}
|
||
|
||
// aggregateHostKeys 汇总若干 host:port 条目的总量、速率与活跃连接数
|
||
func (c *Collector) aggregateHostKeys(keys []string, now time.Time) (uploadTotal, downloadTotal int64, uploadRate, downloadRate float64, connections int) {
|
||
c.mu.RLock()
|
||
defer c.mu.RUnlock()
|
||
|
||
var upSamples, downSamples int64
|
||
var oldest time.Time
|
||
clientSet := map[string]time.Time{}
|
||
|
||
for _, key := range keys {
|
||
st, ok := c.hosts[key]
|
||
if !ok {
|
||
continue
|
||
}
|
||
uploadTotal += st.uploadTotal
|
||
downloadTotal += st.downloadTotal
|
||
cutoff := now.Add(-rateWindow)
|
||
for _, s := range st.samples {
|
||
if s.at.Before(cutoff) {
|
||
continue
|
||
}
|
||
upSamples += s.up
|
||
downSamples += s.down
|
||
if oldest.IsZero() || s.at.Before(oldest) {
|
||
oldest = s.at
|
||
}
|
||
}
|
||
activeCutoff := now.Add(-activeWindow)
|
||
for ip, seen := range st.clients {
|
||
if seen.Before(activeCutoff) {
|
||
continue
|
||
}
|
||
if prev, ok := clientSet[ip]; !ok || seen.After(prev) {
|
||
clientSet[ip] = seen
|
||
}
|
||
}
|
||
}
|
||
|
||
span := now.Sub(oldest).Seconds()
|
||
if span < 0.5 {
|
||
span = 0.5
|
||
}
|
||
if upSamples > 0 {
|
||
uploadRate = float64(upSamples) / span
|
||
}
|
||
if downSamples > 0 {
|
||
downloadRate = float64(downSamples) / span
|
||
}
|
||
connections = len(clientSet)
|
||
return uploadTotal, downloadTotal, uploadRate, downloadRate, connections
|
||
}
|
||
|
||
func (c *Collector) SnapshotForRules(rules []proxy.Rule) []RuleTraffic {
|
||
now := time.Now()
|
||
out := make([]RuleTraffic, 0, len(rules))
|
||
for _, rule := range rules {
|
||
endpoints := rule.Endpoints()
|
||
up, down, upRate, downRate, conns := c.aggregateEndpoints(endpoints, now)
|
||
out = append(out, RuleTraffic{
|
||
RuleID: rule.ID,
|
||
UploadTotal: up,
|
||
DownloadTotal: down,
|
||
UploadRate: upRate,
|
||
DownloadRate: downRate,
|
||
Connections: conns,
|
||
})
|
||
}
|
||
return out
|
||
}
|
||
|
||
// HostTraffic 单个域名的当前速率与总量（指标历史按域名落点）
|
||
type HostTraffic struct {
|
||
Hostname string `json:"hostname"`
|
||
UploadTotal int64 `json:"upload_total"`
|
||
DownloadTotal int64 `json:"download_total"`
|
||
UploadRate float64 `json:"upload_rate"`
|
||
DownloadRate float64 `json:"download_rate"`
|
||
Connections int `json:"connections"`
|
||
}
|
||
|
||
// SnapshotByHost 按域名汇总当前速率：内部键是 host:port，
|
||
// 同一域名的多个端口会合并，与 aggregateEndpoints 的算法保持一致。
|
||
func (c *Collector) SnapshotByHost(hostnames []string) map[string]HostTraffic {
|
||
now := time.Now()
|
||
out := map[string]HostTraffic{}
|
||
for _, host := range hostnames {
|
||
host = strings.ToLower(strings.TrimSpace(host))
|
||
if host == "" {
|
||
continue
|
||
}
|
||
prefix := host + ":"
|
||
c.mu.RLock()
|
||
keys := make([]string, 0, 2)
|
||
for key := range c.hosts {
|
||
if strings.HasPrefix(key, prefix) {
|
||
keys = append(keys, key)
|
||
}
|
||
}
|
||
c.mu.RUnlock()
|
||
if len(keys) == 0 {
|
||
continue
|
||
}
|
||
up, down, upRate, downRate, conns := c.aggregateHostKeys(keys, now)
|
||
out[host] = HostTraffic{
|
||
Hostname: host,
|
||
UploadTotal: up,
|
||
DownloadTotal: down,
|
||
UploadRate: upRate,
|
||
DownloadRate: downRate,
|
||
Connections: conns,
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func (c *Collector) ClientsForRule(rule proxy.Rule) []ClientConn {
|
||
now := time.Now()
|
||
activeCutoff := now.Add(-activeWindow)
|
||
clientSet := map[string]time.Time{}
|
||
|
||
c.mu.RLock()
|
||
for _, ep := range rule.Endpoints() {
|
||
key := EndpointKey(ep.Hostname, ep.Port)
|
||
st, ok := c.hosts[key]
|
||
if !ok {
|
||
continue
|
||
}
|
||
for ip, seen := range st.clients {
|
||
if seen.Before(activeCutoff) {
|
||
continue
|
||
}
|
||
if prev, ok := clientSet[ip]; !ok || seen.After(prev) {
|
||
clientSet[ip] = seen
|
||
}
|
||
}
|
||
}
|
||
c.mu.RUnlock()
|
||
|
||
out := make([]ClientConn, 0, len(clientSet))
|
||
for ip, seen := range clientSet {
|
||
out = append(out, ClientConn{IP: ip, LastSeen: seen.UTC().Format(time.RFC3339)})
|
||
}
|
||
sort.Slice(out, func(i, j int) bool {
|
||
return out[i].LastSeen > out[j].LastSeen
|
||
})
|
||
return out
|
||
}
|