/* Copyright © 2023 Seednode */ package cmd import ( "encoding/gob" "encoding/json" "fmt" "net/http" "os" "os/signal" "sort" "strconv" "strings" "sync" "syscall" "time" "github.com/julienschmidt/httprouter" "github.com/klauspost/compress/zstd" ) type serveStats struct { mutex sync.RWMutex list []string count map[string]uint32 size map[string]string times map[string][]string } type publicServeStats struct { List []string Count map[string]uint32 Size map[string]string Times map[string][]string } type timesServed struct { File string Served uint32 Size string Times []string } func (stats *serveStats) incrementCounter(file string, timestamp time.Time, filesize string) { stats.mutex.Lock() stats.count[file]++ stats.times[file] = append(stats.times[file], timestamp.Format(logDate)) _, exists := stats.size[file] if !exists { stats.size[file] = filesize } if !contains(stats.list, file) { stats.list = append(stats.list, file) } stats.mutex.Unlock() } func (stats *serveStats) Import(source *publicServeStats) { stats.mutex.Lock() copy(stats.list, source.List) for k, v := range source.Count { stats.count[k] = v } for k, v := range source.Size { stats.size[k] = v } for k, v := range source.Times { stats.times[k] = v } stats.mutex.Unlock() } func (source *serveStats) Export() *publicServeStats { stats := &publicServeStats{ List: make([]string, len(source.list)), Count: make(map[string]uint32, len(source.count)), Size: make(map[string]string, len(source.size)), Times: make(map[string][]string, len(source.times)), } source.mutex.RLock() copy(stats.List, source.list) for k, v := range source.count { stats.Count[k] = v } for k, v := range source.size { stats.Size[k] = v } for k, v := range source.times { stats.Times[k] = v } source.mutex.RUnlock() return stats } func (source *serveStats) listFiles(page int) ([]byte, error) { stats := source.Export() sort.SliceStable(stats.List, func(p, q int) bool { return strings.ToLower(stats.List[p]) < strings.ToLower(stats.List[q]) }) var startIndex, stopIndex int if page == -1 { startIndex = 0 stopIndex = len(stats.List) } else { startIndex = ((page - 1) * int(PageLength)) stopIndex = (startIndex + int(PageLength)) } if startIndex > len(stats.List)-1 { return []byte("{}"), nil } if stopIndex > len(stats.List) { stopIndex = len(stats.List) } a := make([]timesServed, (stopIndex - startIndex)) for k, v := range stats.List[startIndex:stopIndex] { a[k] = timesServed{v, stats.Count[v], stats.Size[v], stats.Times[v]} } r, err := json.MarshalIndent(a, "", " ") if err != nil { return []byte{}, err } return r, nil } func (stats *serveStats) ExportFile(path string) error { file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600) if err != nil { return err } defer file.Close() z, err := zstd.NewWriter(file) if err != nil { return err } defer z.Close() enc := gob.NewEncoder(z) err = enc.Encode(stats.Export()) if err != nil { return err } return nil } func (stats *serveStats) ImportFile(path string) error { file, err := os.OpenFile(path, os.O_RDONLY, 0600) if err != nil { return err } defer file.Close() z, err := zstd.NewReader(file) if err != nil { return err } defer z.Close() dec := gob.NewDecoder(z) source := &publicServeStats{ List: []string{}, Count: make(map[string]uint32), Size: make(map[string]string), Times: make(map[string][]string), } err = dec.Decode(source) if err != nil { return err } stats.Import(source) return nil } func serveStatsPage(args []string, stats *serveStats) httprouter.Handle { return func(w http.ResponseWriter, r *http.Request, p httprouter.Params) { startTime := time.Now() page, err := strconv.Atoi(p.ByName("page")) if err != nil || page == 0 { page = -1 } response, err := stats.listFiles(page) if err != nil { fmt.Println(err) serverError(w, r, nil) return } w.Header().Set("Content-Type", "application/json") w.Write(response) if Verbose { fmt.Printf("%s | Served statistics page (%s) to %s in %s\n", startTime.Format(logDate), humanReadableSize(len(response)), realIP(r), time.Since(startTime).Round(time.Microsecond), ) } if StatisticsFile != "" { stats.ExportFile(StatisticsFile) } } } func registerStatsHandlers(mux *httprouter.Router, args []string, stats *serveStats) { if StatisticsFile != "" { stats.ImportFile(StatisticsFile) gracefulShutdown := make(chan os.Signal, 1) signal.Notify(gracefulShutdown, syscall.SIGINT, syscall.SIGTERM) go func() { <-gracefulShutdown stats.ExportFile(StatisticsFile) os.Exit(0) }() mux.GET("/stats", serveStatsPage(args, stats)) if PageLength != 0 { mux.GET("/stats/:page", serveStatsPage(args, stats)) } } }