feat(api, stats): add stats manager and endpoint call tracking

This commit is contained in:
wood chen 2024-10-26 12:39:00 +08:00
parent 89c15abe53
commit fe18d6c828
2 changed files with 119 additions and 0 deletions

22
main.go
View File

@ -11,6 +11,7 @@ import (
"net/url" "net/url"
"os" "os"
"path/filepath" "path/filepath"
"random-api-go/stats"
"strings" "strings"
"sync" "sync"
"time" "time"
@ -31,6 +32,8 @@ var (
rng *rand.Rand rng *rand.Rand
) )
var statsManager *stats.StatsManager
type URLSelector struct { type URLSelector struct {
URLs []string URLs []string
CurrentIndex int CurrentIndex int
@ -80,11 +83,19 @@ func (us *URLSelector) GetRandomURL() string {
return us.URLs[0] return us.URLs[0]
} }
func init() {
// 确保数据目录存在
if err := os.MkdirAll("data", 0755); err != nil {
log.Fatal("Failed to create data directory:", err)
}
}
func main() { func main() {
source := rand.NewSource(time.Now().UnixNano()) source := rand.NewSource(time.Now().UnixNano())
rng = rand.New(source) rng = rand.New(source)
setupLogging() setupLogging()
statsManager = stats.NewStatsManager("data/stats.json")
if err := loadCSVPaths(); err != nil { if err := loadCSVPaths(); err != nil {
log.Fatal("Failed to load CSV paths:", err) log.Fatal("Failed to load CSV paths:", err)
@ -97,6 +108,8 @@ func main() {
// 设置 API 路由 // 设置 API 路由
http.HandleFunc("/pic/", handleAPIRequest) http.HandleFunc("/pic/", handleAPIRequest)
http.HandleFunc("/video/", handleAPIRequest) http.HandleFunc("/video/", handleAPIRequest)
// 添加统计API路由
http.HandleFunc("/stats", handleStats)
log.Printf("Listening on %s...\n", port) log.Printf("Listening on %s...\n", port)
if err := http.ListenAndServe(port, nil); err != nil { if err := http.ListenAndServe(port, nil); err != nil {
@ -252,9 +265,18 @@ func handleAPIRequest(w http.ResponseWriter, r *http.Request) {
randomURL := selector.GetRandomURL() randomURL := selector.GetRandomURL()
// 记录统计
endpoint := fmt.Sprintf("%s/%s", prefix, suffix)
statsManager.IncrementCalls(endpoint)
duration := time.Since(start) duration := time.Since(start)
log.Printf("Request: %s %s from %s - Source: %s - Duration: %v - Redirecting to: %s", log.Printf("Request: %s %s from %s - Source: %s - Duration: %v - Redirecting to: %s",
r.Method, r.URL.Path, realIP, sourceDomain, duration, randomURL) r.Method, r.URL.Path, realIP, sourceDomain, duration, randomURL)
http.Redirect(w, r, randomURL, http.StatusFound) http.Redirect(w, r, randomURL, http.StatusFound)
} }
func handleStats(w http.ResponseWriter, r *http.Request) {
stats := statsManager.GetStats()
json.NewEncoder(w).Encode(stats)
}

97
stats/stats.go Normal file
View File

@ -0,0 +1,97 @@
package stats
import (
"encoding/json"
"os"
"sync"
"time"
)
type EndpointStats struct {
TotalCalls int64 `json:"total_calls"`
TodayCalls int64 `json:"today_calls"`
LastResetDate string `json:"last_reset_date"`
}
type StatsManager struct {
Stats map[string]*EndpointStats `json:"stats"`
mu sync.RWMutex
filepath string
}
func NewStatsManager(filepath string) *StatsManager {
sm := &StatsManager{
Stats: make(map[string]*EndpointStats),
filepath: filepath,
}
sm.LoadStats()
go sm.startDailyReset()
return sm
}
func (sm *StatsManager) IncrementCalls(endpoint string) {
sm.mu.Lock()
defer sm.mu.Unlock()
if _, exists := sm.Stats[endpoint]; !exists {
sm.Stats[endpoint] = &EndpointStats{
LastResetDate: time.Now().Format("2006-01-02"),
}
}
sm.Stats[endpoint].TotalCalls++
sm.Stats[endpoint].TodayCalls++
// 异步保存统计数据
go sm.SaveStats()
}
func (sm *StatsManager) GetStats() map[string]*EndpointStats {
sm.mu.RLock()
defer sm.mu.RUnlock()
return sm.Stats
}
func (sm *StatsManager) SaveStats() error {
sm.mu.RLock()
defer sm.mu.RUnlock()
data, err := json.MarshalIndent(sm, "", " ")
if err != nil {
return err
}
return os.WriteFile(sm.filepath, data, 0644)
}
func (sm *StatsManager) LoadStats() error {
data, err := os.ReadFile(sm.filepath)
if err != nil {
if os.IsNotExist(err) {
return nil
}
return err
}
return json.Unmarshal(data, sm)
}
func (sm *StatsManager) startDailyReset() {
for {
now := time.Now()
next := now.Add(24 * time.Hour)
next = time.Date(next.Year(), next.Month(), next.Day(), 0, 0, 0, 0, next.Location())
duration := next.Sub(now)
time.Sleep(duration)
sm.mu.Lock()
for _, stats := range sm.Stats {
stats.TodayCalls = 0
stats.LastResetDate = time.Now().Format("2006-01-02")
}
sm.mu.Unlock()
sm.SaveStats()
}
}