github.com/yaling888/clash@v1.53.0/tunnel/statistic/manager.go (about)

     1  package statistic
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/samber/lo"
     8  	"go.uber.org/atomic"
     9  )
    10  
    11  var DefaultManager *Manager
    12  
    13  func init() {
    14  	DefaultManager = &Manager{
    15  		uploadTemp:    atomic.NewInt64(0),
    16  		downloadTemp:  atomic.NewInt64(0),
    17  		uploadBlip:    atomic.NewInt64(0),
    18  		downloadBlip:  atomic.NewInt64(0),
    19  		uploadTotal:   atomic.NewInt64(0),
    20  		downloadTotal: atomic.NewInt64(0),
    21  	}
    22  
    23  	go DefaultManager.handle()
    24  }
    25  
    26  type Manager struct {
    27  	connections   sync.Map
    28  	uploadTemp    *atomic.Int64
    29  	downloadTemp  *atomic.Int64
    30  	uploadBlip    *atomic.Int64
    31  	downloadBlip  *atomic.Int64
    32  	uploadTotal   *atomic.Int64
    33  	downloadTotal *atomic.Int64
    34  }
    35  
    36  func (m *Manager) Join(c tracker) {
    37  	m.connections.Store(c.ID(), c)
    38  }
    39  
    40  func (m *Manager) Leave(c tracker) {
    41  	m.connections.Delete(c.ID())
    42  }
    43  
    44  func (m *Manager) KickOut(names ...string) {
    45  	seen := lo.SliceToMap(names, func(item string) (string, struct{}) {
    46  		return item, struct{}{}
    47  	})
    48  
    49  	snapshot := m.Snapshot()
    50  	for _, c := range snapshot.Connections {
    51  		if v, ok := c.(*tcpTracker); ok {
    52  			if _, ok = seen[v.Chains().Last()]; ok {
    53  				_ = v.Close()
    54  			}
    55  		}
    56  	}
    57  }
    58  
    59  func (m *Manager) PushUploaded(size int64) {
    60  	m.uploadTemp.Add(size)
    61  	m.uploadTotal.Add(size)
    62  }
    63  
    64  func (m *Manager) PushDownloaded(size int64) {
    65  	m.downloadTemp.Add(size)
    66  	m.downloadTotal.Add(size)
    67  }
    68  
    69  func (m *Manager) Now() (up int64, down int64) {
    70  	return m.uploadBlip.Load(), m.downloadBlip.Load()
    71  }
    72  
    73  func (m *Manager) Snapshot() *Snapshot {
    74  	connections := []tracker{}
    75  	m.connections.Range(func(key, value any) bool {
    76  		connections = append(connections, value.(tracker))
    77  		return true
    78  	})
    79  
    80  	return &Snapshot{
    81  		UploadTotal:   m.uploadTotal.Load(),
    82  		DownloadTotal: m.downloadTotal.Load(),
    83  		Connections:   connections,
    84  	}
    85  }
    86  
    87  func (m *Manager) ResetStatistic() {
    88  	m.uploadTemp.Store(0)
    89  	m.uploadBlip.Store(0)
    90  	m.uploadTotal.Store(0)
    91  	m.downloadTemp.Store(0)
    92  	m.downloadBlip.Store(0)
    93  	m.downloadTotal.Store(0)
    94  }
    95  
    96  func (m *Manager) Cleanup() {
    97  	snapshot := m.Snapshot()
    98  	for _, c := range snapshot.Connections {
    99  		_ = c.Close()
   100  	}
   101  }
   102  
   103  func (m *Manager) handle() {
   104  	ticker := time.NewTicker(time.Second)
   105  
   106  	for range ticker.C {
   107  		m.uploadBlip.Store(m.uploadTemp.Load())
   108  		m.uploadTemp.Store(0)
   109  		m.downloadBlip.Store(m.downloadTemp.Load())
   110  		m.downloadTemp.Store(0)
   111  	}
   112  }
   113  
   114  type Snapshot struct {
   115  	DownloadTotal int64     `json:"downloadTotal"`
   116  	UploadTotal   int64     `json:"uploadTotal"`
   117  	Connections   []tracker `json:"connections"`
   118  }