Added stats endpoint

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> [path]... [flags]
+  roulette <path> [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)

cmd/files.go

@@ -5,6 +5,7 @@ Copyright © 2023 Seednode <seednode@seedno.de>
 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,

cmd/root.go

@@ -5,7 +5,7 @@ Copyright © 2023 Seednode <seednode@seedno.de>
 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)
+			log.Fatal(err)
-			os.Exit(1)
 		}
 	},
 }
@@ -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")

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)

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)