JCS-pub/client/internal/accessstat/access_stat.go

package accessstat

import (
	"sync"
	"time"

	"gitlink.org.cn/cloudream/common/pkgs/async"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/db"
	clitypes "gitlink.org.cn/cloudream/jcs-pub/client/types"
)

type AccessStatEventChan = async.UnboundChannel[AccessStatEvent]

type AccessStatEvent interface {
	IsAccessStatEvent() bool
}

type ExitEvent struct {
	AccessStatEvent
	Err error
}

type AccessStat struct {
	cfg   Config
	done  chan any
	stats map[entryKey]float64
	lock  sync.Mutex
	db    *db.DB
}

type entryKey struct {
	objID   clitypes.ObjectID
	pkgID   clitypes.PackageID
	spaceID clitypes.UserSpaceID
}

func NewAccessStat(cfg Config, db *db.DB) *AccessStat {
	return &AccessStat{
		cfg:   cfg,
		done:  make(chan any),
		stats: make(map[entryKey]float64),
		db:    db,
	}
}

func (p *AccessStat) AddAccessCounter(objID clitypes.ObjectID, pkgID clitypes.PackageID, spaceID clitypes.UserSpaceID, value float64) {
	p.lock.Lock()
	defer p.lock.Unlock()

	key := entryKey{
		objID:   objID,
		pkgID:   pkgID,
		spaceID: spaceID,
	}
	p.stats[key] += value
}

func (p *AccessStat) Start() *AccessStatEventChan {
	ch := async.NewUnboundChannel[AccessStatEvent]()

	go func() {
		ticker := time.NewTicker(p.cfg.ReportInterval)
		defer ticker.Stop()

		for {
			select {
			case <-ticker.C:
			case <-p.done:
				ch.Send(ExitEvent{})
				return
			}

			p.lock.Lock()
			st := p.stats
			p.stats = make(map[entryKey]float64)
			p.lock.Unlock()

			if len(st) == 0 {
				continue
			}

			var entries []db.AddAccessStatEntry
			for k, v := range st {
				entries = append(entries, db.AddAccessStatEntry{
					ObjectID:    k.objID,
					PackageID:   k.pkgID,
					UserSpaceID: k.spaceID,
					Counter:     v,
				})
			}
			err := db.DoTx10(p.db, p.db.Package().BatchAddPackageAccessStat, entries)

			if err != nil {
				logger.Errorf("add all package access stat counter: %v", err)

				p.lock.Lock()
				for k, v := range st {
					p.stats[k] += v
				}
				p.lock.Unlock()
				continue
			}
		}
	}()
	return ch
}

func (p *AccessStat) Stop() {
	close(p.done)
}