Skip to content

Commit

Permalink
feat(api, stats): add stats manager and endpoint call tracking
Browse files Browse the repository at this point in the history
  • Loading branch information
woodchen-ink committed Oct 26, 2024
1 parent 89c15ab commit fe18d6c
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 0 deletions.
22 changes: 22 additions & 0 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"net/url"
"os"
"path/filepath"
"random-api-go/stats"
"strings"
"sync"
"time"
Expand All @@ -31,6 +32,8 @@ var (
rng *rand.Rand
)

var statsManager *stats.StatsManager

type URLSelector struct {
URLs []string
CurrentIndex int
Expand Down Expand Up @@ -80,11 +83,19 @@ func (us *URLSelector) GetRandomURL() string {
return us.URLs[0]
}

func init() {
// 确保数据目录存在
if err := os.MkdirAll("data", 0755); err != nil {
log.Fatal("Failed to create data directory:", err)
}
}

func main() {
source := rand.NewSource(time.Now().UnixNano())
rng = rand.New(source)

setupLogging()
statsManager = stats.NewStatsManager("data/stats.json")

if err := loadCSVPaths(); err != nil {
log.Fatal("Failed to load CSV paths:", err)
Expand All @@ -97,6 +108,8 @@ func main() {
// 设置 API 路由
http.HandleFunc("/pic/", handleAPIRequest)
http.HandleFunc("/video/", handleAPIRequest)
// 添加统计API路由
http.HandleFunc("/stats", handleStats)

log.Printf("Listening on %s...\n", port)
if err := http.ListenAndServe(port, nil); err != nil {
Expand Down Expand Up @@ -252,9 +265,18 @@ func handleAPIRequest(w http.ResponseWriter, r *http.Request) {

randomURL := selector.GetRandomURL()

// 记录统计
endpoint := fmt.Sprintf("%s/%s", prefix, suffix)
statsManager.IncrementCalls(endpoint)

duration := time.Since(start)
log.Printf("Request: %s %s from %s - Source: %s - Duration: %v - Redirecting to: %s",
r.Method, r.URL.Path, realIP, sourceDomain, duration, randomURL)

http.Redirect(w, r, randomURL, http.StatusFound)
}

func handleStats(w http.ResponseWriter, r *http.Request) {
stats := statsManager.GetStats()
json.NewEncoder(w).Encode(stats)
}
97 changes: 97 additions & 0 deletions stats/stats.go
Original file line number Diff line number Diff line change
@@ -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()
}
}

0 comments on commit fe18d6c

Please sign in to comment.