github.com/metacubex/mihomo@v1.18.5/tunnel/statistic/manager.go (about)

     1  package statistic
     2  
     3  import (
     4  	"os"
     5  	"time"
     6  
     7  	"github.com/metacubex/mihomo/common/atomic"
     8  
     9  	"github.com/puzpuzpuz/xsync/v3"
    10  	"github.com/shirou/gopsutil/v3/process"
    11  )
    12  
    13  var DefaultManager *Manager
    14  
    15  func init() {
    16  	DefaultManager = &Manager{
    17  		connections:   xsync.NewMapOf[string, Tracker](),
    18  		uploadTemp:    atomic.NewInt64(0),
    19  		downloadTemp:  atomic.NewInt64(0),
    20  		uploadBlip:    atomic.NewInt64(0),
    21  		downloadBlip:  atomic.NewInt64(0),
    22  		uploadTotal:   atomic.NewInt64(0),
    23  		downloadTotal: atomic.NewInt64(0),
    24  		process:       &process.Process{Pid: int32(os.Getpid())},
    25  	}
    26  
    27  	go DefaultManager.handle()
    28  }
    29  
    30  type Manager struct {
    31  	connections   *xsync.MapOf[string, Tracker]
    32  	uploadTemp    atomic.Int64
    33  	downloadTemp  atomic.Int64
    34  	uploadBlip    atomic.Int64
    35  	downloadBlip  atomic.Int64
    36  	uploadTotal   atomic.Int64
    37  	downloadTotal atomic.Int64
    38  	process       *process.Process
    39  	memory        uint64
    40  }
    41  
    42  func (m *Manager) Join(c Tracker) {
    43  	m.connections.Store(c.ID(), c)
    44  }
    45  
    46  func (m *Manager) Leave(c Tracker) {
    47  	m.connections.Delete(c.ID())
    48  }
    49  
    50  func (m *Manager) Get(id string) (c Tracker) {
    51  	if value, ok := m.connections.Load(id); ok {
    52  		c = value
    53  	}
    54  	return
    55  }
    56  
    57  func (m *Manager) Range(f func(c Tracker) bool) {
    58  	m.connections.Range(func(key string, value Tracker) bool {
    59  		return f(value)
    60  	})
    61  }
    62  
    63  func (m *Manager) PushUploaded(size int64) {
    64  	m.uploadTemp.Add(size)
    65  	m.uploadTotal.Add(size)
    66  }
    67  
    68  func (m *Manager) PushDownloaded(size int64) {
    69  	m.downloadTemp.Add(size)
    70  	m.downloadTotal.Add(size)
    71  }
    72  
    73  func (m *Manager) Now() (up int64, down int64) {
    74  	return m.uploadBlip.Load(), m.downloadBlip.Load()
    75  }
    76  
    77  func (m *Manager) Memory() uint64 {
    78  	m.updateMemory()
    79  	return m.memory
    80  }
    81  
    82  func (m *Manager) Snapshot() *Snapshot {
    83  	var connections []*TrackerInfo
    84  	m.Range(func(c Tracker) bool {
    85  		connections = append(connections, c.Info())
    86  		return true
    87  	})
    88  	return &Snapshot{
    89  		UploadTotal:   m.uploadTotal.Load(),
    90  		DownloadTotal: m.downloadTotal.Load(),
    91  		Connections:   connections,
    92  		Memory:        m.memory,
    93  	}
    94  }
    95  
    96  func (m *Manager) updateMemory() {
    97  	stat, err := m.process.MemoryInfo()
    98  	if err != nil {
    99  		return
   100  	}
   101  	m.memory = stat.RSS
   102  }
   103  
   104  func (m *Manager) ResetStatistic() {
   105  	m.uploadTemp.Store(0)
   106  	m.uploadBlip.Store(0)
   107  	m.uploadTotal.Store(0)
   108  	m.downloadTemp.Store(0)
   109  	m.downloadBlip.Store(0)
   110  	m.downloadTotal.Store(0)
   111  }
   112  
   113  func (m *Manager) handle() {
   114  	ticker := time.NewTicker(time.Second)
   115  
   116  	for range ticker.C {
   117  		m.uploadBlip.Store(m.uploadTemp.Load())
   118  		m.uploadTemp.Store(0)
   119  		m.downloadBlip.Store(m.downloadTemp.Load())
   120  		m.downloadTemp.Store(0)
   121  	}
   122  }
   123  
   124  type Snapshot struct {
   125  	DownloadTotal int64          `json:"downloadTotal"`
   126  	UploadTotal   int64          `json:"uploadTotal"`
   127  	Connections   []*TrackerInfo `json:"connections"`
   128  	Memory        uint64         `json:"memory"`
   129  }