cmd/btrtrc/txstats.go

   1 package main
   2
   3 import (
   4         "fmt"
   5         "log"
   6         "os"
   7         "strconv"
   8         "sync"
   9         "time"
  10
  11         "github.com/anacrolix/torrent"
  12         "github.com/anacrolix/torrent/metainfo"
  13 )
  14
  15 const TxExt = ".tx"
  16
  17 var (
  18         TxStats  = map[metainfo.Hash]int64{}
  19         TxStatsM sync.Mutex
  20 )
  21
  22 func txStatsLoad(h metainfo.Hash) {
  23         pth := h.HexString() + TxExt
  24         data, err := os.ReadFile(pth)
  25         if err != nil {
  26                 return
  27         }
  28         v, err := strconv.ParseInt(string(data[:len(data)-1]), 10, 64)
  29         if err != nil {
  30                 log.Println("ParseInt:", pth, err)
  31                 return
  32         }
  33         TxStatsM.Lock()
  34         TxStats[h] = v
  35         TxStatsM.Unlock()
  36 }
  37
  38 func txStatsDel(h metainfo.Hash) {
  39         TxStatsM.Lock()
  40         delete(TxStats, h)
  41         TxStatsM.Unlock()
  42 }
  43
  44 func txStatsDump(t *torrent.Torrent) {
  45         stats := t.Stats()
  46         TxStatsM.Lock()
  47         s := stats.BytesWrittenData.Int64() + TxStats[t.InfoHash()]
  48         pth := t.InfoHash().HexString() + TxExt
  49         if err := os.WriteFile(pth, []byte(fmt.Sprintf("%d\n", s)), 0666); err != nil {
  50                 log.Println("WriteFile:", pth, err)
  51         }
  52         TxStatsM.Unlock()
  53 }
  54
  55 func txStatsDumpAll(c *torrent.Client) {
  56         for _, t := range c.Torrents() {
  57                 if t.Info() != nil {
  58                         txStatsDump(t)
  59                 }
  60         }
  61 }
  62
  63 func txStatsDumper(c *torrent.Client) {
  64         tick := time.Tick(10 * time.Second)
  65         for {
  66                 txStatsDumpAll(c)
  67                 select {
  68                 case <-Cancel:
  69                         Jobs.Done()
  70                         return
  71                 case <-tick:
  72                 }
  73         }
  74 }