package metrics import ( "log" "strconv" "strings" "time" "github.com/shirou/gopsutil/v3/cpu" "github.com/shirou/gopsutil/v3/load" "github.com/shirou/gopsutil/v3/mem" "github.com/wormhole/wormhole/internal/bridge" "github.com/wormhole/wormhole/internal/cache" "github.com/wormhole/wormhole/internal/store" ) type Collector struct { store *store.Store cache *cache.Manager bridge *bridge.Bridge stop chan struct{} } func NewCollector(st *store.Store, c *cache.Manager, b *bridge.Bridge) *Collector { return &Collector{ store: st, cache: c, bridge: b, stop: make(chan struct{}), } } func (c *Collector) Start(flushInterval time.Duration) { go func() { ticker := time.NewTicker(flushInterval) defer ticker.Stop() for { select { case <-ticker.C: c.flushFlow() case <-c.stop: return } } }() } func (c *Collector) Stop() { close(c.stop) } func (c *Collector) flushFlow() { deltas := c.cache.DrainFlow() for key, delta := range deltas { parts := strings.SplitN(key, ":", 2) if len(parts) != 2 { continue } resourceType := parts[0] resourceID, err := strconv.ParseUint(parts[1], 10, 64) if err != nil { continue } if err := c.store.UpsertFlowStat(resourceType, uint(resourceID), delta.Inlet, delta.Export); err != nil { log.Printf("[metrics] flush flow stat: %v", err) } switch resourceType { case "client": _ = c.store.AddClientFlow(uint(resourceID), delta.Inlet, delta.Export) case "tunnel": _ = c.store.AddTunnelFlow(uint(resourceID), delta.Inlet, delta.Export) case "host": _ = c.store.AddHostFlow(uint(resourceID), delta.Inlet, delta.Export) } } } func (c *Collector) RecordRequestIP(resourceType string, resourceID uint, visitorIP, directIP string) { if err := c.store.RecordRequestIP(resourceType, resourceID, visitorIP, directIP); err != nil { log.Printf("[metrics] record request ip: %v", err) } } func (c *Collector) AddFlow(resourceType string, resourceID uint, clientID uint, inlet, export int64) { c.cache.AddFlow(resourceType, resourceID, inlet, export) if clientID > 0 { c.cache.AddFlow("client", clientID, inlet, export) } } func (c *Collector) Dashboard() (map[string]interface{}, error) { stats, err := c.store.DashboardStats() if err != nil { return nil, err } stats["client_online"] = c.bridge.OnlineCount() cpuPercent, _ := cpu.Percent(0, false) memInfo, _ := mem.VirtualMemory() loadInfo, _ := load.Avg() var cpuVal float64 if len(cpuPercent) > 0 { cpuVal = cpuPercent[0] } topClients, _ := c.store.TopClientsByFlow(10) topTunnels, _ := c.store.TopTunnelsByFlow(10) topHosts, _ := c.store.TopHostsByFlow(10) topIPs, _ := c.store.TopRequestIPs(10) stats["cpu_percent"] = cpuVal stats["mem_percent"] = memInfo.UsedPercent stats["load1"] = loadInfo.Load1 stats["top_clients"] = topClients stats["top_tunnels"] = topTunnels stats["top_hosts"] = topHosts stats["top_ips"] = topIPs return stats, nil }