diff --git a/handlers/url_stats.go b/handlers/url_stats.go new file mode 100644 index 0000000..cd08233 --- /dev/null +++ b/handlers/url_stats.go @@ -0,0 +1,22 @@ +package handlers + +import ( + "encoding/json" + "net/http" + "random-api-go/services" +) + +func HandleURLStats(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + w.Header().Set("Content-Type", "application/json") + stats := services.GetURLCounts() + + if err := json.NewEncoder(w).Encode(stats); err != nil { + http.Error(w, "Error encoding response", http.StatusInternalServerError) + return + } +} diff --git a/main.go b/main.go index 1467d69..5db7963 100644 --- a/main.go +++ b/main.go @@ -9,6 +9,7 @@ import ( "random-api-go/config" "random-api-go/handlers" "random-api-go/logging" + "random-api-go/services" "random-api-go/stats" "syscall" @@ -36,6 +37,11 @@ func main() { log.Fatal("Failed to initialize handlers:", err) } + // 初始化加载所有CSV内容 + if err := services.InitializeCSVService(); err != nil { + log.Fatal("Failed to initialize CSV Service:", err) + } + // 设置路由 setupRoutes() @@ -64,4 +70,6 @@ func setupRoutes() { http.HandleFunc("/pic/", handlers.HandleAPIRequest) http.HandleFunc("/video/", handlers.HandleAPIRequest) http.HandleFunc("/stats", handlers.HandleStats) + // 添加URL统计接口 + http.HandleFunc("/urlstats", handlers.HandleURLStats) } diff --git a/services/csv_service.go b/services/csv_service.go index ca5143a..1976c43 100644 --- a/services/csv_service.go +++ b/services/csv_service.go @@ -23,7 +23,32 @@ var ( // InitializeCSVService 初始化CSV服务 func InitializeCSVService() error { - return LoadCSVPaths() + // 加载url.json + if err := LoadCSVPaths(); err != nil { + return fmt.Errorf("failed to load CSV paths: %v", err) + } + + // 预加载所有CSV内容 + Mu.RLock() + defer Mu.RUnlock() + + for prefix, suffixMap := range CSVPathsCache { + for suffix, csvPath := range suffixMap { + selector, err := GetCSVContent(csvPath) + if err != nil { + log.Printf("Warning: Failed to load CSV content for %s/%s: %v", prefix, suffix, err) + continue + } + + // 更新URL计数 + endpoint := fmt.Sprintf("%s/%s", prefix, suffix) + UpdateURLCount(endpoint, csvPath, len(selector.URLs)) + + log.Printf("Loaded %d URLs for endpoint: %s/%s", len(selector.URLs), prefix, suffix) + } + } + + return nil } func LoadCSVPaths() error { diff --git a/services/url_counter.go b/services/url_counter.go new file mode 100644 index 0000000..13e9aaf --- /dev/null +++ b/services/url_counter.go @@ -0,0 +1,43 @@ +package services + +import "sync" + +type URLStats struct { + TotalURLs int `json:"total_urls"` +} + +var ( + URLCounts = make(map[string]struct { + stats URLStats + csvPath string // 内部使用,不会输出到JSON + }) + urlMu sync.RWMutex +) + +// 更新URL计数 +func UpdateURLCount(endpoint string, csvPath string, count int) { + urlMu.Lock() + defer urlMu.Unlock() + URLCounts[endpoint] = struct { + stats URLStats + csvPath string + }{ + stats: URLStats{ + TotalURLs: count, + }, + csvPath: csvPath, + } +} + +// 获取URL计数 +func GetURLCounts() map[string]URLStats { + urlMu.RLock() + defer urlMu.RUnlock() + + // 返回一个只包含统计信息的副本 + counts := make(map[string]URLStats) + for k, v := range URLCounts { + counts[k] = v.stats + } + return counts +}