github.com/sagernet/sing-box@v1.2.7/experimental/clashapi/trafficontrol/manager.go (about)

     1  package trafficontrol
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/sagernet/sing-box/experimental/clashapi/compatible"
     7  	"github.com/sagernet/sing/common/atomic"
     8  )
     9  
    10  type Manager struct {
    11  	uploadTemp    atomic.Int64
    12  	downloadTemp  atomic.Int64
    13  	uploadBlip    atomic.Int64
    14  	downloadBlip  atomic.Int64
    15  	uploadTotal   atomic.Int64
    16  	downloadTotal atomic.Int64
    17  
    18  	connections compatible.Map[string, tracker]
    19  	ticker      *time.Ticker
    20  	done        chan struct{}
    21  }
    22  
    23  func NewManager() *Manager {
    24  	manager := &Manager{
    25  		ticker: time.NewTicker(time.Second),
    26  		done:   make(chan struct{}),
    27  	}
    28  	go manager.handle()
    29  	return manager
    30  }
    31  
    32  func (m *Manager) Join(c tracker) {
    33  	m.connections.Store(c.ID(), c)
    34  }
    35  
    36  func (m *Manager) Leave(c tracker) {
    37  	m.connections.Delete(c.ID())
    38  }
    39  
    40  func (m *Manager) PushUploaded(size int64) {
    41  	m.uploadTemp.Add(size)
    42  	m.uploadTotal.Add(size)
    43  }
    44  
    45  func (m *Manager) PushDownloaded(size int64) {
    46  	m.downloadTemp.Add(size)
    47  	m.downloadTotal.Add(size)
    48  }
    49  
    50  func (m *Manager) Now() (up int64, down int64) {
    51  	return m.uploadBlip.Load(), m.downloadBlip.Load()
    52  }
    53  
    54  func (m *Manager) Snapshot() *Snapshot {
    55  	var connections []tracker
    56  	m.connections.Range(func(_ string, value tracker) bool {
    57  		connections = append(connections, value)
    58  		return true
    59  	})
    60  
    61  	return &Snapshot{
    62  		UploadTotal:   m.uploadTotal.Load(),
    63  		DownloadTotal: m.downloadTotal.Load(),
    64  		Connections:   connections,
    65  	}
    66  }
    67  
    68  func (m *Manager) ResetStatistic() {
    69  	m.uploadTemp.Store(0)
    70  	m.uploadBlip.Store(0)
    71  	m.uploadTotal.Store(0)
    72  	m.downloadTemp.Store(0)
    73  	m.downloadBlip.Store(0)
    74  	m.downloadTotal.Store(0)
    75  }
    76  
    77  func (m *Manager) handle() {
    78  	var uploadTemp int64
    79  	var downloadTemp int64
    80  	for {
    81  		select {
    82  		case <-m.done:
    83  			return
    84  		case <-m.ticker.C:
    85  		}
    86  		uploadTemp = m.uploadTemp.Swap(0)
    87  		downloadTemp = m.downloadTemp.Swap(0)
    88  		m.uploadBlip.Store(uploadTemp)
    89  		m.downloadBlip.Store(downloadTemp)
    90  	}
    91  }
    92  
    93  func (m *Manager) Close() error {
    94  	m.ticker.Stop()
    95  	close(m.done)
    96  	return nil
    97  }
    98  
    99  type Snapshot struct {
   100  	DownloadTotal int64     `json:"downloadTotal"`
   101  	UploadTotal   int64     `json:"uploadTotal"`
   102  	Connections   []tracker `json:"connections"`
   103  }