From 537633a03028ee51807ac49114f9f7ab08b3a99f Mon Sep 17 00:00:00 2001 From: Seednode Date: Wed, 8 Feb 2023 07:50:40 -0600 Subject: [PATCH] Added persistent stats file --- cmd/files.go | 1 - cmd/root.go | 18 +++--- cmd/version.go | 2 +- cmd/web.go | 166 +++++++++++++++++++++++++++++++++++++++++++------ 4 files changed, 157 insertions(+), 30 deletions(-) diff --git a/cmd/files.go b/cmd/files.go index c8ad504..8219b98 100644 --- a/cmd/files.go +++ b/cmd/files.go @@ -318,7 +318,6 @@ func splitPath(path string, Regexes *Regexes) (*Path, error) { p.base = split[0][1] p.number, err = strconv.Atoi(split[0][2]) - if err != nil { return &Path{}, err } diff --git a/cmd/root.go b/cmd/root.go index 3a3815f..6a089fa 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -12,14 +12,15 @@ import ( ) var ( - cache bool - cacheFile string - filtering bool - port uint16 - recursive bool - sorting bool - statistics bool - verbose bool + cache bool + cacheFile string + filtering bool + port uint16 + recursive bool + sorting bool + statistics bool + statisticsFile string + verbose bool rootCmd = &cobra.Command{ Use: "roulette [path]...", @@ -52,6 +53,7 @@ func init() { rootCmd.Flags().BoolVarP(&recursive, "recursive", "r", false, "recurse into subdirectories") rootCmd.Flags().BoolVarP(&sorting, "sort", "s", false, "enable sorting") rootCmd.Flags().BoolVar(&statistics, "stats", false, "expose stats endpoint") + rootCmd.Flags().StringVar(&statisticsFile, "stats-file", "", "path to optional persistent stats file") rootCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "log accessed files to stdout") rootCmd.Flags().SetInterspersed(true) } diff --git a/cmd/version.go b/cmd/version.go index 4d92996..5166e59 100644 --- a/cmd/version.go +++ b/cmd/version.go @@ -10,7 +10,7 @@ import ( "github.com/spf13/cobra" ) -var Version = "0.36.2" +var Version = "0.37" func init() { rootCmd.AddCommand(versionCmd) diff --git a/cmd/web.go b/cmd/web.go index e87a475..a2fd475 100644 --- a/cmd/web.go +++ b/cmd/web.go @@ -10,10 +10,12 @@ import ( "errors" "fmt" "io" + "log" "math/rand" "net/http" "net/url" "os" + "os/signal" "path/filepath" "regexp" "runtime" @@ -21,6 +23,7 @@ import ( "strconv" "strings" "sync" + "syscall" "time" "github.com/klauspost/compress/zstd" @@ -122,16 +125,18 @@ func (i *Index) Export(path string) error { enc := gob.NewEncoder(z) + i.mutex.RLock() + enc.Encode(&i.list) + i.mutex.RUnlock() + return nil } func (i *Index) Import(path string) error { file, err := os.OpenFile(path, os.O_RDONLY, 0600) - if os.IsNotExist(err) { - return ErrIndexNotExist - } else if err != nil { + if err != nil { return err } defer file.Close() @@ -144,7 +149,12 @@ func (i *Index) Import(path string) error { dec := gob.NewDecoder(z) + i.mutex.Lock() + err = dec.Decode(&i.list) + + i.mutex.Unlock() + if err != nil { return err } @@ -160,6 +170,13 @@ type ServeStats struct { times map[string][]string } +type exportedServeStats struct { + List []string + Count map[string]uint64 + Size map[string]string + Times map[string][]string +} + func (s *ServeStats) incrementCounter(image string, timestamp time.Time, filesize string) { s.mutex.Lock() @@ -179,27 +196,66 @@ func (s *ServeStats) incrementCounter(image string, timestamp time.Time, filesiz s.mutex.Unlock() } -func (s *ServeStats) ListImages() ([]byte, error) { +func (s *ServeStats) toExported() *exportedServeStats { + stats := &exportedServeStats{ + List: []string{}, + Count: make(map[string]uint64), + Size: make(map[string]string), + Times: make(map[string][]string), + } + s.mutex.RLock() - sortedList := &ServeStats{ - mutex: sync.RWMutex{}, - list: s.list, - count: s.count, - size: s.size, - times: s.times, + stats.List = append(stats.List, s.list...) + + for k, v := range s.count { + stats.Count[k] = v + } + + for k, v := range s.size { + stats.Size[k] = v + } + + for k, v := range s.times { + stats.Times[k] = v } s.mutex.RUnlock() - sort.SliceStable(sortedList.list, func(p, q int) bool { - return sortedList.list[p] < sortedList.list[q] + return stats +} + +func (s *ServeStats) toImported(stats *exportedServeStats) { + s.mutex.Lock() + + s.list = append(s.list, stats.List...) + + for k, v := range stats.Count { + s.count[k] = v + } + + for k, v := range stats.Size { + s.size[k] = v + } + + for k, v := range stats.Times { + s.times[k] = v + } + + s.mutex.Unlock() +} + +func (s *ServeStats) ListImages() ([]byte, error) { + stats := s.toExported() + + sort.SliceStable(stats.List, func(p, q int) bool { + return stats.List[p] < stats.List[q] }) a := []timesServed{} - for _, image := range sortedList.list { - a = append(a, timesServed{image, sortedList.count[image], sortedList.size[image], sortedList.times[image]}) + for _, image := range stats.List { + a = append(a, timesServed{image, stats.Count[image], stats.Size[image], stats.Times[image]}) } r, err := json.MarshalIndent(a, "", " ") @@ -210,6 +266,63 @@ func (s *ServeStats) ListImages() ([]byte, error) { return r, nil } +func (s *ServeStats) Export(path string) error { + file, err := os.OpenFile(path, os.O_RDWR|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) + + stats := s.toExported() + + err = enc.Encode(&stats) + if err != nil { + log.Fatal(err) + } + + return nil +} + +func (s *ServeStats) Import(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) + + stats := &exportedServeStats{ + List: []string{}, + Count: make(map[string]uint64), + Size: make(map[string]string), + Times: make(map[string][]string), + } + + err = dec.Decode(stats) + if err != nil { + return err + } + + s.toImported(stats) + + return nil +} + type timesServed struct { File string Served uint64 @@ -522,6 +635,10 @@ func serveStatsHandler(args []string, stats *ServeStats) http.HandlerFunc { time.Since(startTime).Round(time.Microsecond), ) } + + if statisticsFile != "" { + stats.Export(statisticsFile) + } } } @@ -655,13 +772,9 @@ func ServePage(args []string) error { skipIndex := false if cacheFile != "" { - err = index.Import(cacheFile) - switch err { - case ErrIndexNotExist: - case nil: + err := index.Import(cacheFile) + if err == nil { skipIndex = true - default: - return err } } @@ -680,6 +793,19 @@ func ServePage(args []string) error { times: make(map[string][]string), } + if statistics && statisticsFile != "" { + stats.Import(statisticsFile) + + gracefulShutdown := make(chan os.Signal, 1) + signal.Notify(gracefulShutdown, syscall.SIGINT, syscall.SIGTERM) + + go func() { + <-gracefulShutdown + stats.Export(statisticsFile) + os.Exit(0) + }() + } + http.Handle("/", serveHtmlHandler(paths, Regexes, index)) http.Handle(Prefix+"/", http.StripPrefix(Prefix, serveStaticFileHandler(paths, stats))) http.HandleFunc("/favicon.ico", doNothing)