github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/components/statistics/cache.go (about)

     1  package statistics
     2  
     3  import (
     4  	"encoding/binary"
     5  	"sync/atomic"
     6  
     7  	"github.com/Asutorufa/yuhaiin/pkg/utils/cache"
     8  )
     9  
    10  var (
    11  	DownloadKey = []byte{'D', 'O', 'W', 'N', 'L', 'O', 'A', 'D'}
    12  	UploadKey   = []byte{'U', 'P', 'L', 'O', 'A', 'D'}
    13  
    14  	SyncThreshold int64 = 1024 * 1024 * 50 // bytes
    15  )
    16  
    17  type Cache struct {
    18  	download atomic.Uint64
    19  	upload   atomic.Uint64
    20  
    21  	notSyncDownload atomic.Int64
    22  	notSyncUpload   atomic.Int64
    23  
    24  	cache *cache.Cache
    25  }
    26  
    27  func NewCache(cache *cache.Cache) *Cache {
    28  	c := &Cache{
    29  		cache: cache,
    30  	}
    31  
    32  	if download := cache.Get(DownloadKey); download != nil {
    33  		c.download.Store(binary.BigEndian.Uint64(download))
    34  	}
    35  
    36  	if upload := cache.Get(UploadKey); upload != nil {
    37  		c.upload.Store(binary.BigEndian.Uint64(upload))
    38  	}
    39  
    40  	return c
    41  }
    42  
    43  func (c *Cache) AddDownload(d uint64) {
    44  	c.download.Add(d)
    45  
    46  	z := c.notSyncDownload.Add(int64(d))
    47  	if z >= SyncThreshold {
    48  		c.cache.Put(DownloadKey, binary.BigEndian.AppendUint64(nil, c.download.Load()))
    49  		c.notSyncDownload.Add(-z)
    50  	}
    51  }
    52  
    53  func (c *Cache) LoadDownload() uint64 { return c.download.Load() }
    54  
    55  func (c *Cache) AddUpload(d uint64) {
    56  	c.upload.Add(d)
    57  
    58  	z := c.notSyncUpload.Add(int64(d))
    59  	if z >= SyncThreshold {
    60  		c.cache.Put(UploadKey, binary.BigEndian.AppendUint64(nil, c.upload.Load()))
    61  		c.notSyncUpload.Add(-z)
    62  	}
    63  }
    64  
    65  func (c *Cache) LoadUpload() uint64 { return c.upload.Load() }
    66  
    67  func (c *Cache) Close() {
    68  	c.cache.Put(DownloadKey, binary.BigEndian.AppendUint64(nil, c.download.Load()))
    69  	c.cache.Put(UploadKey, binary.BigEndian.AppendUint64(nil, c.upload.Load()))
    70  }