From 1a4496b788037624050d4e9bc64664ef6ef894ec Mon Sep 17 00:00:00 2001 From: Seednode Date: Thu, 19 Jan 2023 12:07:15 -0600 Subject: [PATCH] Added stats endpoint --- README.md | 9 +++++- cmd/files.go | 74 +++++++++++++++++++++++++++++++++++++++++--------- cmd/root.go | 7 +++-- cmd/version.go | 2 +- cmd/web.go | 44 +++++++++++++++++++++++++++--- 5 files changed, 114 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 06403bd..17ec7ab 100644 --- a/README.md +++ b/README.md @@ -74,10 +74,16 @@ If any `include=`/`exclude=` filters are specified in a given request, the cache The cache can be regenerated any time by accessing the `/clear_cache` endpoint. +## Debug + +If the `-d|--debug` flag is passed, an additional endpoint, `/stats`, is registered. + +When accessed, this endpoint returns a JSON document listing every file served, along with the number of times it has been served. + ## Usage output ``` Usage: - roulette [path]... [flags] + roulette [path2]... [flags] roulette [command] Available Commands: @@ -87,6 +93,7 @@ Available Commands: Flags: -c, --cache only scan directories once, at startup (or when filters are applied) + -d, --debug store list of files served and number of times they were served -f, --filter enable filtering via query parameters -h, --help help for roulette -p, --port uint16 port to listen on (default 8080) diff --git a/cmd/files.go b/cmd/files.go index 5109bfd..b52422d 100644 --- a/cmd/files.go +++ b/cmd/files.go @@ -5,6 +5,7 @@ Copyright © 2023 Seednode package cmd import ( + "encoding/json" "errors" "fmt" "image" @@ -43,37 +44,75 @@ type Files struct { List map[string][]string } -type Stats struct { +type ScanStats struct { FilesMatched uint64 FilesSkipped uint64 DirectoriesMatched uint64 } -func (s *Stats) GetFilesTotal() uint64 { +type TimesServed struct { + File string + Count uint64 +} + +type ServeStats struct { + ImagesServed uint64 + ImageList []string + ImageCount map[string]uint64 +} + +func (s *ServeStats) GetFilesTotal() uint64 { + return atomic.LoadUint64(&s.ImagesServed) +} + +func (s *ServeStats) IncrementCounter(image string) { + s.ImageCount[image] += 1 + + if !contains(s.ImageList, image) { + s.ImageList = append(s.ImageList, image) + } +} + +func (s *ServeStats) ListImages() ([]byte, error) { + a := []TimesServed{} + + for _, image := range s.ImageList { + a = append(a, TimesServed{image, s.ImageCount[image]}) + } + + r, err := json.Marshal(a) + if err != nil { + return []byte{}, err + } + + return r, nil +} + +func (s *ScanStats) GetFilesTotal() uint64 { return atomic.LoadUint64(&s.FilesMatched) + atomic.LoadUint64(&s.FilesSkipped) } -func (s *Stats) IncrementFilesMatched() { +func (s *ScanStats) IncrementFilesMatched() { atomic.AddUint64(&s.FilesMatched, 1) } -func (s *Stats) GetFilesMatched() uint64 { +func (s *ScanStats) GetFilesMatched() uint64 { return atomic.LoadUint64(&s.FilesMatched) } -func (s *Stats) IncrementFilesSkipped() { +func (s *ScanStats) IncrementFilesSkipped() { atomic.AddUint64(&s.FilesSkipped, 1) } -func (s *Stats) GetFilesSkipped() uint64 { +func (s *ScanStats) GetFilesSkipped() uint64 { return atomic.LoadUint64(&s.FilesSkipped) } -func (s *Stats) IncrementDirectoriesMatched() { +func (s *ScanStats) IncrementDirectoriesMatched() { atomic.AddUint64(&s.DirectoriesMatched, 1) } -func (s *Stats) GetDirectoriesMatched() uint64 { +func (s *ScanStats) GetDirectoriesMatched() uint64 { return atomic.LoadUint64(&s.DirectoriesMatched) } @@ -91,6 +130,15 @@ func (p *Path) Decrement() { p.Number = p.Number - 1 } +func contains(s []string, e string) bool { + for _, a := range s { + if a == e { + return true + } + } + return false +} + func humanReadableSize(bytes int) string { const unit = 1000 @@ -138,7 +186,7 @@ func preparePath(path string) string { return path } -func appendPath(directory, path string, files *Files, stats *Stats) error { +func appendPath(directory, path string, files *Files, stats *ScanStats) error { // If caching, only check image types once, during the initial scan, to speed up future pickFile() calls if Cache { image, err := isImage(path) @@ -160,7 +208,7 @@ func appendPath(directory, path string, files *Files, stats *Stats) error { return nil } -func appendPaths(path string, files *Files, filters *Filters, stats *Stats) error { +func appendPaths(path string, files *Files, filters *Filters, stats *ScanStats) error { absolutePath, err := filepath.Abs(path) if err != nil { return err @@ -372,7 +420,7 @@ func isImage(path string) (bool, error) { return filetype.IsImage(head), nil } -func getFiles(path string, files *Files, filters *Filters, stats *Stats, concurrency *Concurrency) error { +func getFiles(path string, files *Files, filters *Filters, stats *ScanStats, concurrency *Concurrency) error { var wg sync.WaitGroup err := filepath.WalkDir(path, func(p string, info os.DirEntry, err error) error { @@ -414,7 +462,7 @@ func getFiles(path string, files *Files, filters *Filters, stats *Stats, concurr return nil } -func getFileList(paths []string, files *Files, filters *Filters, stats *Stats, concurrency *Concurrency) { +func getFileList(paths []string, files *Files, filters *Filters, stats *ScanStats, concurrency *Concurrency) { var wg sync.WaitGroup for i := 0; i < len(paths); i++ { @@ -489,7 +537,7 @@ func pickFile(args []string, filters *Filters, sort string, fileCache *[]string) List: make(map[string][]string), } - stats := &Stats{ + stats := &ScanStats{ FilesMatched: 0, FilesSkipped: 0, DirectoriesMatched: 0, diff --git a/cmd/root.go b/cmd/root.go index 89f21a7..9503f31 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -5,7 +5,7 @@ Copyright © 2023 Seednode package cmd import ( - "fmt" + "log" "os" "github.com/spf13/cobra" @@ -25,6 +25,7 @@ type Concurrency struct { } var Cache bool +var Debug bool var Filter bool var Port uint16 var Recursive bool @@ -38,8 +39,7 @@ var rootCmd = &cobra.Command{ Run: func(cmd *cobra.Command, args []string) { err := ServePage(args) if err != nil { - fmt.Println(err) - os.Exit(1) + log.Fatal(err) } }, } @@ -53,6 +53,7 @@ func Execute() { func init() { rootCmd.Flags().BoolVarP(&Cache, "cache", "c", false, "only scan directories once, at startup (or when filters are applied)") + rootCmd.Flags().BoolVarP(&Debug, "debug", "d", false, "store list of files served and number of times they were served") rootCmd.Flags().BoolVarP(&Filter, "filter", "f", false, "enable filtering via query parameters") rootCmd.Flags().Uint16VarP(&Port, "port", "p", 8080, "port to listen on") rootCmd.Flags().BoolVarP(&Recursive, "recursive", "r", false, "recurse into subdirectories") diff --git a/cmd/version.go b/cmd/version.go index 786ff9b..158a31a 100644 --- a/cmd/version.go +++ b/cmd/version.go @@ -10,7 +10,7 @@ import ( "github.com/spf13/cobra" ) -var Version = "0.30.2" +var Version = "0.31.0" func init() { rootCmd.AddCommand(versionCmd) diff --git a/cmd/web.go b/cmd/web.go index 1dea9df..31bfa46 100644 --- a/cmd/web.go +++ b/cmd/web.go @@ -277,7 +277,7 @@ func serveHtml(w http.ResponseWriter, r *http.Request, filePath string, dimensio return nil } -func serveStaticFile(w http.ResponseWriter, r *http.Request, paths []string) error { +func serveStaticFile(w http.ResponseWriter, r *http.Request, paths []string, stats *ServeStats) error { prefixedFilePath, err := stripQueryParams(r.URL.Path) if err != nil { return err @@ -307,6 +307,8 @@ func serveStaticFile(w http.ResponseWriter, r *http.Request, paths []string) err startTime := time.Now() + stats.IncrementCounter(filePath) + buf, err := os.ReadFile(filePath) if err != nil { return err @@ -353,9 +355,34 @@ func serveCacheClearHandler(args []string, fileCache *[]string) http.HandlerFunc } } -func serveStaticFileHandler(paths []string) http.HandlerFunc { +func serveStats(args []string, fileCache *[]string) error { + filters := &Filters{} + + fileCache = &[]string{} + + fmt.Printf("%v | Preparing image cache...\n", time.Now().Format(LogDate)) + _, err := pickFile(args, filters, "", fileCache) + + return err +} + +func serveStatsHandler(args []string, stats *ServeStats) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { - err := serveStaticFile(w, r, paths) + w.WriteHeader(http.StatusOK) + w.Header().Set("Content-Type", "application/json") + + response, err := stats.ListImages() + if err != nil { + log.Fatal(err) + } + + w.Write([]byte(response)) + } +} + +func serveStaticFileHandler(paths []string, stats *ServeStats) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + err := serveStaticFile(w, r, paths, stats) if err != nil { log.Fatal(err) } @@ -476,9 +503,18 @@ func ServePage(args []string) error { } } + stats := &ServeStats{ + ImagesServed: 0, + ImageList: []string{}, + ImageCount: make(map[string]uint64), + } + http.Handle("/", serveHtmlHandler(paths, regexes, fileCache)) http.Handle("/clear_cache", serveCacheClearHandler(args, fileCache)) - http.Handle(Prefix+"/", http.StripPrefix(Prefix, serveStaticFileHandler(paths))) + + http.Handle("/stats", serveStatsHandler(args, stats)) + + http.Handle(Prefix+"/", http.StripPrefix(Prefix, serveStaticFileHandler(paths, stats))) http.HandleFunc("/favicon.ico", doNothing) err = http.ListenAndServe(":"+strconv.FormatInt(int64(Port), 10), nil)