1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
|
package main
import (
"context"
"database/sql"
"flag"
"os"
"os/signal"
"time"
_ "github.com/mattn/go-sqlite3"
"github.com/sirupsen/logrus"
"google.golang.org/grpc"
"go.rikki.moe/v2stat/command"
)
var (
flagDatabase = flag.String("db", "v2stat.db", "Path to SQLite database")
flagServer = flag.String("server", "127.0.0.1:8080", "V2Ray API server address")
flagLogLevel = flag.String("log-level", "info", "Log level (debug, info, warn, error, fatal, panic)")
)
// V2Stat holds references to the logger, database connection, and gRPC client.
type V2Stat struct {
logger *logrus.Logger
db *sql.DB
stat command.StatsServiceClient
}
func main() {
flag.Parse()
// Initialize logger
level, err := logrus.ParseLevel(*flagLogLevel)
if err != nil {
logrus.Fatalf("Invalid log level: %v", err)
}
logger := logrus.New()
logger.SetLevel(level)
// Dial gRPC server
conn, err := grpc.NewClient(*flagServer, grpc.WithInsecure())
if err != nil {
logger.Fatalf("Failed to dial gRPC server: %v", err)
}
defer conn.Close()
statClient := command.NewStatsServiceClient(conn)
// Open SQLite database
db, err := sql.Open("sqlite3", *flagDatabase)
if err != nil {
logger.Fatalf("Failed to open database: %v", err)
}
defer db.Close()
// Create main struct
v2stat := &V2Stat{
logger: logger,
db: db,
stat: statClient,
}
// Initialize database schema
if err := v2stat.InitDB(); err != nil {
logger.Fatalf("Failed to initialize database: %v", err)
}
// For graceful shutdown, create a context that cancels on SIGINT/SIGTERM
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, os.Interrupt, os.Kill)
// Optional: Query stats once with a reset (as in your original code)
if _, err := v2stat.stat.QueryStats(ctx, &command.QueryStatsRequest{Reset_: true}); err != nil {
logger.Errorf("Failed to query stats: %v", err)
}
// Wait until the top of the next hour
now := time.Now()
if now.Minute() != 0 || now.Second() != 0 {
nextHour := now.Truncate(time.Hour).Add(time.Hour)
subDuration := nextHour.Sub(now)
logger.Infof("Waiting for %s to start recording stats", subDuration)
timer := time.NewTimer(subDuration)
select {
case <-timer.C:
case <-sigCh:
logger.Info("Received shutdown signal, exiting.")
return
}
}
// Start a ticker for every hour
ticker := time.NewTicker(1 * time.Hour)
defer ticker.Stop()
// Main loop
for {
logger.Info("Recording stats...")
if err := v2stat.RecordNow(ctx); err != nil {
logger.Errorf("Failed to record stats: %v", err)
}
// Wait for next ticker or shutdown signal
select {
case <-ticker.C:
// just continue the loop and record again
case <-sigCh:
logger.Info("Received shutdown signal, exiting.")
return
case <-ctx.Done():
logger.Info("Context canceled, exiting.")
return
}
}
}
|
