Files
vps-control/control-plane/backend/poller.go
DizelBadCoder c4dc19cd41
Some checks failed
CI / containers (push) Has been cancelled
CI / agent (push) Has been cancelled
CI / backend (push) Has been cancelled
CI / frontend (push) Has been cancelled
git init
2026-09-19 19:31:36 +03:00

255 lines
7.5 KiB
Go

package main
import (
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"net/url"
"regexp"
"sync"
"time"
)
var bearerSecret = regexp.MustCompile(`(?i)(authorization\s*[:=]\s*bearer\s+)[^\s,;&]+`)
var namedSecret = regexp.MustCompile(`(?i)((?:password|passwd|token|secret|api[_-]?key)["']?\s*[:=]\s*)(?:"[^"]*"|'[^']*'|[^\s,;&]+)`)
type Poller struct {
store *Store
agent *AgentClient
metricRetention time.Duration
logRetention time.Duration
}
func (poller *Poller) Run(ctx context.Context) {
pollTicker := time.NewTicker(time.Minute)
cleanupTicker := time.NewTicker(time.Hour)
defer pollTicker.Stop()
defer cleanupTicker.Stop()
poller.poll(ctx)
for {
select {
case <-ctx.Done():
return
case <-pollTicker.C:
poller.poll(ctx)
case <-cleanupTicker.C:
if _, err := poller.store.pool.Exec(ctx, "DELETE FROM metric_samples WHERE sampled_at < $1", time.Now().Add(-poller.metricRetention)); err != nil {
log.Printf("metric retention cleanup: %v", err)
}
if _, err := poller.store.pool.Exec(ctx, "DELETE FROM agent_log_entries WHERE event_at < $1", time.Now().Add(-poller.logRetention)); err != nil {
log.Printf("log retention cleanup: %v", err)
}
_, _ = poller.store.pool.Exec(ctx, "DELETE FROM sessions WHERE expires_at < now()")
}
}
}
func (poller *Poller) poll(ctx context.Context) {
servers, err := poller.store.Servers(ctx)
if err != nil {
log.Printf("poll server list: %v", err)
return
}
jobs := make(chan Server)
var wait sync.WaitGroup
workerCount := 4
if len(servers) < workerCount {
workerCount = len(servers)
}
for worker := 0; worker < workerCount; worker++ {
wait.Add(1)
go func() {
defer wait.Done()
for server := range jobs {
poller.pollServer(ctx, server)
}
}()
}
for _, server := range servers {
select {
case jobs <- server:
case <-ctx.Done():
close(jobs)
wait.Wait()
return
}
}
close(jobs)
wait.Wait()
}
func (poller *Poller) pollServer(ctx context.Context, server Server) {
statusCode, body, err := poller.agent.Do(ctx, server, http.MethodGet, "/v1/system/status", nil)
if err != nil || statusCode != http.StatusOK {
_, _ = poller.store.pool.Exec(ctx, "UPDATE servers SET status='offline',updated_at=now() WHERE id=$1", server.ID)
return
}
var status struct {
Hostname string `json:"hostname"`
CPU float64 `json:"cpu_usage_percent"`
MemoryTotal uint64 `json:"memory_total_bytes"`
MemoryUsed uint64 `json:"memory_used_bytes"`
Filesystems []struct {
Total uint64 `json:"total_bytes"`
Used uint64 `json:"used_bytes"`
} `json:"filesystems"`
}
if json.Unmarshal(body, &status) != nil {
return
}
memory := percent(status.MemoryUsed, status.MemoryTotal)
var diskTotal, diskUsed uint64
for _, filesystem := range status.Filesystems {
diskTotal += filesystem.Total
diskUsed += filesystem.Used
}
disk := percent(diskUsed, diskTotal)
var rx, tx int64
if code, payload, callErr := poller.agent.Do(ctx, server, http.MethodGet, "/v1/network/interfaces", nil); callErr == nil && code == 200 {
var interfaces []struct {
RX int64 `json:"rx_bytes"`
TX int64 `json:"tx_bytes"`
}
if json.Unmarshal(payload, &interfaces) == nil {
for _, item := range interfaces {
rx += item.RX
tx += item.TX
}
}
}
connections := 0
if code, payload, callErr := poller.agent.Do(ctx, server, http.MethodGet, "/v1/network/connections", nil); callErr == nil && code == 200 {
var items []json.RawMessage
if json.Unmarshal(payload, &items) == nil {
connections = len(items)
}
}
txDB, err := poller.store.pool.Begin(ctx)
if err != nil {
return
}
defer txDB.Rollback(ctx)
_, err = txDB.Exec(ctx, "UPDATE servers SET status='online',hostname=$2,last_seen=now(),updated_at=now() WHERE id=$1", server.ID, status.Hostname)
if err == nil {
_, err = txDB.Exec(ctx, `INSERT INTO metric_samples(server_id,cpu,memory,disk,rx,tx,connections,bans)
VALUES($1,$2,$3,$4,$5,$6,$7,0)`, server.ID, status.CPU, memory, disk, rx, tx, connections)
}
if err == nil {
err = txDB.Commit(ctx)
}
if err != nil {
log.Printf("store metric for %s: %v", server.ID, err)
}
poller.pollLogs(ctx, server)
}
func (poller *Poller) pollLogs(ctx context.Context, server Server) {
var cursor string
var since time.Time
err := poller.store.pool.QueryRow(ctx, `SELECT
COALESCE((SELECT journal_cursor FROM agent_log_entries WHERE server_id=$1 ORDER BY event_at DESC,ingested_at DESC LIMIT 1),''),
COALESCE((SELECT max(event_at) FROM agent_log_entries WHERE server_id=$1),now() - interval '2 minutes')`, server.ID).
Scan(&cursor, &since)
if err != nil {
log.Printf("log cursor for %s: %v", server.ID, err)
return
}
fallbackUsed := false
for page := 0; page < 5; page++ {
path := fmt.Sprintf("/v1/logs?since=%d&limit=200", since.Add(-time.Second).Unix())
if cursor != "" {
path = "/v1/logs?cursor=" + url.QueryEscape(cursor) + "&limit=200"
}
status, payload, callErr := poller.agent.Do(ctx, server, http.MethodGet, path, nil)
if callErr != nil || status == http.StatusForbidden || status == http.StatusNotFound {
return
}
if status != http.StatusOK {
if cursor != "" && !fallbackUsed {
cursor = ""
fallbackUsed = true
page--
continue
}
log.Printf("log poll for %s returned %d", server.ID, status)
return
}
var response struct {
Entries []struct {
Cursor string `json:"cursor"`
Timestamp uint64 `json:"timestamp_unix_micros"`
Unit string `json:"unit"`
Priority int `json:"priority"`
Message string `json:"message"`
BootID *string `json:"boot_id"`
PID *int `json:"pid"`
} `json:"entries"`
Truncated bool `json:"truncated"`
}
if err = json.Unmarshal(payload, &response); err != nil {
log.Printf("decode logs for %s: %v", server.ID, err)
return
}
if len(response.Entries) > 200 {
log.Printf("agent %s exceeded log entry bound", server.ID)
return
}
tx, beginErr := poller.store.pool.Begin(ctx)
if beginErr != nil {
return
}
now := time.Now()
for _, entry := range response.Entries {
if len(entry.Cursor) == 0 || len(entry.Cursor) > 1024 || len(entry.Unit) > 128 ||
len(entry.Message) > 65_536 || entry.Priority < 0 || entry.Priority > 7 {
continue
}
eventAt := time.UnixMicro(int64(entry.Timestamp))
if eventAt.After(now.Add(5*time.Minute)) || eventAt.Before(now.Add(-31*24*time.Hour)) {
continue
}
_, err = tx.Exec(ctx, `INSERT INTO agent_log_entries
(server_id,journal_cursor,event_at,unit,priority,message,boot_id,pid)
VALUES($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT DO NOTHING`, server.ID, entry.Cursor,
eventAt, entry.Unit, entry.Priority, truncateUTF8(redactLogMessage(entry.Message), 65_536), entry.BootID, entry.PID)
if err != nil {
_ = tx.Rollback(ctx)
return
}
}
if err = tx.Commit(ctx); err != nil {
log.Printf("store logs for %s: %v", server.ID, err)
return
}
if len(response.Entries) == 0 || !response.Truncated {
return
}
cursor = response.Entries[len(response.Entries)-1].Cursor
}
log.Printf("agent log backlog for %s exceeded 1000 entries in one poll", server.ID)
}
func redactLogMessage(message string) string {
message = bearerSecret.ReplaceAllString(message, "${1}[REDACTED]")
return namedSecret.ReplaceAllString(message, "${1}[REDACTED]")
}
func truncateUTF8(value string, maximum int) string {
if len(value) <= maximum {
return value
}
for maximum > 0 && value[maximum]&0xc0 == 0x80 {
maximum--
}
return value[:maximum]
}
func percent(used, total uint64) float64 {
if total == 0 {
return 0
}
return float64(used) * 100 / float64(total)
}